Feat: add DeepSeek V4 Flash prefill context parallelism - #889
Conversation
📝 WalkthroughWalkthroughChangesDeepSeek V4 Flash attention paths
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 13
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/prefill_indexer_compressor.py (1)
347-393: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAdd an explicit dependency from
prefill_idx_c4_scale_scattertocache_write_tid.
scale_scratchis written inside theprefill_idx_c4_cache_writeSPMD task at line 368. Theprefill_idx_c4_scale_scatterblock at line 382 readsscale_scratchbut declares nodeps. The producer task idcache_write_tidis already in scope and is only used later at line 425 for the completion signal.If the scheduler does not derive this read-after-write edge automatically, the scatter can publish stale or uninitialized dequantization scales into
idx_kv_scale. That corrupts the C8 indexer cache while the INT8 rows stay correct, which makes the defect hard to trace.🔒️ Proposed fix to order the scale scatter after the cache write
- with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_idx_c4_scale_scatter") as scale_scatter_tid: + with pl.at( + level=pl.Level.CORE_GROUP, + name_hint="prefill_idx_c4_scale_scatter", + deps=[cache_write_tid], + ) as scale_scatter_tid:🤖 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/prefill_indexer_compressor.py` around lines 347 - 393, Add an explicit dependency on cache_write_tid to the prefill_idx_c4_scale_scatter task declaration so it waits for the producer that writes scale_scratch. Preserve the existing scale scatter logic and dependency behavior for all other tasks.
🧹 Nitpick comments (14)
models/deepseek/v4-flash/prefill_csa.py (1)
251-259: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMark the unused indexer outputs as intentional.
Ruff RUF059 reports
idx_kv_cache_outandidx_kv_scale_outas unused. Ifprefill_indexerupdatesidx_kv_cacheandidx_kv_scalein place, use underscore-prefixed bindings.Suggested binding change
- idx_kv_cache_out, idx_kv_scale_out, idx_score_unused, cmp_topk_indices = prefill_indexer( + _idx_kv_cache_out, _idx_kv_scale_out, idx_score_unused, cmp_topk_indices = prefill_indexer(🤖 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/prefill_csa.py` around lines 251 - 259, Update the assignment receiving outputs from prefill_indexer so the intentionally unused idx_kv_cache_out and idx_kv_scale_out bindings use underscore-prefixed names, while preserving the remaining outputs and call behavior.Source: Linters/SAST tools
models/deepseek/v4-flash/prefill_swa_cp.py (2)
502-510: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe loop variable
tail_startshadows the module-leveltail_startfunction.Line 502 binds
tail_startto a scalar insideprefill_cp_swa. Line 142 definestail_start(seg_start, seg_len)at module scope. The function is not called later in this kernel, so the shadowing is currently harmless.Rename the local to
tail_start0, matchingtail_offset0inprefill_hca_cp.pyline 646. That prevents a confusing failure if this kernel later needs the helper.🤖 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/prefill_swa_cp.py` around lines 502 - 510, Rename the local scalar `tail_start` in `prefill_cp_swa` to `tail_start0`, and update its use when computing `tail_offset`; preserve the existing tail-copy behavior while avoiding shadowing the module-level `tail_start` helper.
642-644: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign
build_tensor_specswith the sibling CP modules and add thecp_sizeguard.This builder returns
(specs, ctx)at line 704 and the__main__block installsgolden_prefill_cp_swa._ctxexternally at line 854.prefill_csa_cp.pyline 1007 andprefill_hca_cp.pyline 1546 instead set_ctxinside the builder and return only the spec list. A shared harness that callsbuild_tensor_specsuniformly across these modules will break on this one.This builder also omits the
cp_sizeguard.prefill_hca_cp.pylines 1309-1312 andprefill_csa_cp.pylines 639-642 raise when the runtimecp_sizediffers from the import-timeCP_SIZE. Here,build_metadata(args.cp)would build metadata for one CP size while the compiled kernel shapes use another.Set
_ctxinsidebuild_tensor_specs, return onlyspecs, and raise whencp_size != CP_SIZE.♻️ Proposed change
def build_tensor_specs(cp_size: int = CP_SIZE): + if cp_size != CP_SIZE: + raise ValueError( + f"runtime cp_size={cp_size} does not match static CP_SIZE={CP_SIZE}" + ) meta, ctx = build_metadata(cp_size)specs.append(TensorSpec("x_out", list(x.shape), torch.float32, is_output=True)) - return specs, ctx + golden_prefill_cp_swa._ctx = ctx + return specs- specs, ctx = build_tensor_specs(args.cp) - golden_prefill_cp_swa._ctx = ctx result = run_jit( fn=prefill_cp_swa_test, - specs=specs, + specs=build_tensor_specs(args.cp), golden_fn=golden_prefill_cp_swa,Also applies to: 836-854
🤖 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/prefill_swa_cp.py` around lines 642 - 644, Update build_tensor_specs to reject any cp_size that differs from the import-time CP_SIZE before building metadata, set golden_prefill_cp_swa._ctx within the builder, and return only specs. Adjust the __main__ block to use the builder’s new return contract without externally installing _ctx.models/deepseek/v4-flash/prefill_hca_cp.py (2)
647-662: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCollapse the no-op branch in
cp_hca_tail_assemble.Both arms of the
if tail_offset < TAIL_ROWSbranch compute the samesourceexpression. The branch has no effect.The layout is already correct: tiles of one part are contiguous, so
part * MAX_SEGMENT_TILES * TAIL_ROWS + tail_offsetaddresses tile 0 fortail_offset < TAIL_ROWSand tile 1 above it. The branch reads as an unfinished edit and invites a wrong "fix" later.The equivalent code in
prefill_swa_cp.pylines 506-509 has the same redundancy in a different form (TAIL_ROWS + tail_offset - TAIL_ROWS). Simplify both.♻️ Proposed simplification
for row in pl.range(TAIL_ROWS): tail_offset = tail_offset0 + row if tail_offset < total: - if tail_offset < TAIL_ROWS: - source = ( - part * MAX_SEGMENT_TILES * TAIL_ROWS - + tail_offset - ) - else: - source = ( - part * MAX_SEGMENT_TILES * TAIL_ROWS - + tail_offset - ) + # Tiles of one part are contiguous, so the offset indexes directly. + source = part * MAX_SEGMENT_TILES * TAIL_ROWS + tail_offset destination = part * TAIL_ROWS + row🤖 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/prefill_hca_cp.py` around lines 647 - 662, In cp_hca_tail_assemble, remove the redundant if/else around source and assign it directly using part * MAX_SEGMENT_TILES * TAIL_ROWS + tail_offset. Apply the same simplification to the corresponding tail assembly logic in prefill_swa_cp.py, preserving the existing contiguous tile layout and indexing behavior.
1313-1314: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffTwo metadata builders produce overlapping fixture data.
build_hca_metadataat line 322 and_build_raw_attention_metadataat line 169 both recomputeprefix,span,lengths,starts,owners, andquery_positionsfrom the same inputs.device_valuesat lines 1501-1523 then mixes the two sources:query_positionscomes fromraw_metadata, whilesegment_starts_t,cmp_indices, andsegment_tail_positionscome frommetadata.The two
query_positionstensors also differ in their inactive fill:_build_raw_attention_metadatauses0at line 177 andbuild_hca_metadatauses-1at line 335. The kernel only reads active rows, so today this is benign, but the split makes the fixture contract hard to audit.Derive the shared segment geometry once and pass it to both builders.
🤖 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/prefill_hca_cp.py` around lines 1313 - 1314, Derive the shared segment geometry once before the metadata construction in the fixture setup, then pass that geometry into both build_hca_metadata and _build_raw_attention_metadata instead of recomputing prefix, span, lengths, starts, owners, and query_positions independently. Update both builder signatures and their callers, and ensure device_values uses fields from this single consistent geometry source.models/deepseek/v4-flash/prefill_zigzag_cp.py (1)
33-43: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExport
_parse_static_intinstead of copying it intoprefill_swa_cp.py.
prefill_swa_cp.pylines 75-82 define an identical_parse_static_intand then derive its ownCP_SIZEfromCP_DEFAULT, which it already imports from this module. Two independent argv scans can disagree if the parsing rule ever changes.Import this helper in
prefill_swa_cp.py, or importCP_SIZEdirectly since both modules derive the same value from the same flag.🤖 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/prefill_zigzag_cp.py` around lines 33 - 43, Export and reuse the existing _parse_static_int helper from prefill_zigzag_cp.py in prefill_swa_cp.py, removing the duplicate implementation and deriving CP_SIZE through the shared parser (or importing CP_SIZE directly). Ensure both modules use one consistent --cp argument parsing rule.models/deepseek/v4-flash/prefill_csa_cp.py (1)
271-289: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift
_build_metadata_tensorscomputes a large unused result, and it hides the only capacity check.
build_tensor_specsat line 643 unpacksmetadata, ctxand then reads only four entries:compress_state_block_table,inner_compress_state_block_table,cmp_block_table, andidx_block_table. All four come from_build_block_tables. Ruff also reports thatctxis never used.Two consequences:
- Lines 292-483 build
boundary_positions,seed_*,main_slot_mapping,idx_slot_mapping,final_*_state_mapping,dense_cmp_prefix, anddense_idx_prefixthat no spec consumes. Lines 896-955 then recompute the equivalent leaf slot mappings with the same_lower_rowrule. The two derivations can drift apart, and only the second one reaches the kernel.validate_cp_indexer_capacityat line 289 runs only as a side effect of this otherwise-unused function. If someone removes the dead branch, the capacity guard disappears silently.Reduce
_build_metadata_tensorsto the tables and context thatbuild_tensor_specsactually consumes, and callvalidate_cp_indexer_capacityfrombuild_tensor_specsdirectly.Also applies to: 643-643
🤖 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/prefill_csa_cp.py` around lines 271 - 289, Reduce _build_metadata_tensors to return only the metadata tables and context entries consumed by build_tensor_specs, removing the unused boundary, seed, slot-mapping, state-mapping, and dense-prefix calculations. Move validate_cp_indexer_capacity out of _build_metadata_tensors and invoke it directly from build_tensor_specs, preserving the existing capacity validation behavior independently of metadata construction.Source: Linters/SAST tools
models/deepseek/v4-flash/prefill_exchange_cp.py (1)
478-623: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift
_prefill_cp_sparse_stagestill duplicates_cp_swa_stage_sources.This helper and
_cp_swa_stage_sourcesinprefill_swa_cp.pylines 294-379 implement the same persistent/predecessor/current overlay lowering. The gather loop, the overlay index derivation, and the bias build are line-for-line equivalent. This helper adds the compressed-index branch at lines 564-577 and the valid-mask build at lines 603-621; the SWA copy omits both.The PR objective states that exchange logic is consolidated into
prefill_exchange_cp.py, so the SWA copy is an outlier. The SWA path can call_prefill_cp_sparse_stagewithIDX_TOPK-widecmp_indicesfilled with-1, which makes the compressed branch inert and produces the same valid mask.This is deferable, but it prevents the two overlay lowerings from diverging.
🤖 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/prefill_exchange_cp.py` around lines 478 - 623, Refactor the SWA path to reuse _prefill_cp_sparse_stage instead of maintaining the duplicate _cp_swa_stage_sources implementation in prefill_swa_cp.py. Supply IDX_TOPK-wide cmp_indices initialized to -1 and the required compressed-stage buffers so the compressed branch remains inert, while preserving the existing SWA outputs and overlay behavior.models/deepseek/v4-flash/prefill_indexer.py (1)
379-482: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftConsider extracting the shared query-preparation pipeline.
Lines 379-482 of
_prefill_indexer_cp_score_topkrepeat the body ofprefill_indexerlines 148-253 almost verbatim: theqr_projINT8 GEMM and dequantization, the interleaved RoPE table materialization, the swap-index build, the per-token RoPE rotation, the Hadamard multiply, and the per-row INT8 quantization. Only thename_hintstrings differ.Extract one
@pl.jit.inlinehelper that returnsqr_hadamard_i8andqr_hadamard_scale_dq, and give it aname_prefixargument for the task hints. That keeps the two scoring paths from drifting when the quantization math changes.This is deferable if you prefer to land the CP path first, but track it: the same block already exists in three copies across the CP modules.
🤖 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/prefill_indexer.py` around lines 379 - 482, Extract the duplicated query-preparation pipeline from `_prefill_indexer_cp_score_topk` and `prefill_indexer` into one `@pl.jit.inline` helper that returns `qr_hadamard_i8` and `qr_hadamard_scale_dq`. Move the shared projection/dequantization, RoPE preparation and rotation, Hadamard multiplication, and per-row INT8 quantization into the helper, using a `name_prefix` parameter to generate task-specific `name_hint` values, then replace each duplicated block with the helper call.models/deepseek/v4-flash/decode_sparse_attn_csa.py (4)
52-57: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the duplicate
MAX_SEQ_LENassignment.Line 39 already sets
MAX_SEQ_LEN = M.max_position_embeddings. Line 54 repeats the same assignment inside the CSA masking block, where the preceding comment describes compressed-slot masking. The duplicate suggests a different value is intended here.♻️ Proposed cleanup
# CSA compressed-slot masking (folded in from the CSA orchestrator): raw indexer # topk -> per-token bound floor((pos + 1) / COMPRESS_RATIO). -MAX_SEQ_LEN = M.max_position_embeddings INDEXER_SCORE_LEN = MAX_SEQ_LEN // 4🤖 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/decode_sparse_attn_csa.py` around lines 52 - 57, Remove the duplicate MAX_SEQ_LEN assignment from the CSA compressed-slot masking block, while retaining the existing assignment near line 39. Keep INDEXER_SCORE_LEN and the remaining CSA constants using the shared MAX_SEQ_LEN value.
690-710: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the loop variable to avoid shadowing the masked tensor.
rawis bound at Line 670 to the masked indexer tensor. Both loops here rebindrawto a plainint. The tensor is not read after Line 673, so the result is correct, but the reuse makes the masking block harder to follow.♻️ Proposed rename
- for raw in window_swa_indices[t].tolist(): - slot = int(raw) + for win_slot_raw in window_swa_indices[t].tolist(): + slot = int(win_slot_raw) @@ - for raw in cmp_sparse_indices[t].tolist(): - if raw < 0: + for cmp_slot_raw in cmp_sparse_indices[t].tolist(): + if cmp_slot_raw < 0: kv_rows.append(torch.zeros(HEAD_DIM, dtype=ori_kv.dtype)) valid.append(False) continue - cmp_slot = int(raw) + cmp_slot = int(cmp_slot_raw)🤖 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/decode_sparse_attn_csa.py` around lines 690 - 710, Rename the loop variable raw in both index-iteration loops to distinct descriptive names, updating each corresponding reference while preserving the existing index handling and masking behavior.
542-545: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueHoist the
T_PADzero-fill out of the token loop.This block sits inside
for qt in pl.pipeline(0, T, QUANT_TOKEN_TILE, ...). It writes the same constant zeros to the same padding rows on every token tile, so it repeatsT // QUANT_TOKEN_TILEtimes. The result is correct, but the work is redundant. The padding rows do not depend onqt.♻️ Proposed refactor
o_r_i8_pad = pl.assemble(o_r_i8_pad, oq_i8, [qt, col_g]) - if T_PAD > T: - zero_half = pl.full([T_PAD - T, O_LORA], dtype=pl.FP16, value=0.0) - zero_i8 = pl.cast(zero_half, target_type=pl.INT8, mode="trunc") - o_r_i8_pad = pl.assemble(o_r_i8_pad, zero_i8, [T, col_g]) + if T_PAD > T: + zero_half = pl.full([T_PAD - T, O_LORA], dtype=pl.FP16, value=0.0) + zero_i8 = pl.cast(zero_half, target_type=pl.INT8, mode="trunc") + o_r_i8_pad = pl.assemble(o_r_i8_pad, zero_i8, [T, col_g])🤖 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/decode_sparse_attn_csa.py` around lines 542 - 545, Move the T_PAD padding initialization out of the for qt token-loop and perform it once before the loop, preserving the existing zero_half, zero_i8, and o_r_i8_pad assembly behavior for T_PAD > T. Keep token-dependent processing inside the loop and ensure the padding rows remain unchanged for every tile.
316-337: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider bulk-run detection for the window gather.
This loop issues one
pl.gather_rowper sparse-K row, 128 rows per work item. The sibling kernels avoid that for the window prefix.decode_sparse_attn_swa.sparse_attn_swaanddecode_sparse_attn_hca.sparse_attn_hcaprobe the first and last slot of aGATHER_RUNsub-tile, and copy the whole run as one bulk transfer when the endpoints areGATHER_RUN - 1apart. Window slots are contiguous in the common case, so the same probe applies here.This changes scheduling only, not results. Treat it as deferrable if the current profile is acceptable.
🤖 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/decode_sparse_attn_csa.py` around lines 316 - 337, Optimize the window-prefix path in the loop around qk_kv by detecting contiguous GATHER_RUN sub-tiles, probing the first and last window slots as in decode_sparse_attn_swa.sparse_attn_swa and decode_sparse_attn_hca.sparse_attn_hca, and using one bulk gather when their separation is GATHER_RUN - 1. Preserve the existing per-row gather and fallback behavior for non-contiguous or invalid slots, without changing results for the comparison-sparse path.models/deepseek/v4-flash/decode_swa.py (1)
417-427: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueCache the SWA metadata to avoid recomputing the nested Python loop.
init_swa_indicesandinit_swa_lenseach callinit_swa_metadata, which runsswa_indices_and_lens. That helper loops overT * WINelements with per-element.item()calls. It also rebuildsinit_start_posandinit_block_tableon every call. The values are deterministic, so the results stay consistent, but the fixture pays the cost twice.♻️ Proposed refactor
- def init_swa_metadata(): - return swa_indices_and_lens( - position_ids_from_starts(init_start_pos(), seq=S), - init_block_table(), - block_size=BLOCK_SIZE, - window=WIN, - ) + swa_metadata_cache = {} + + def init_swa_metadata(): + if "value" not in swa_metadata_cache: + swa_metadata_cache["value"] = swa_indices_and_lens( + position_ids_from_starts(init_start_pos(), seq=S), + init_block_table(), + block_size=BLOCK_SIZE, + window=WIN, + ) + return swa_metadata_cache["value"]🤖 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/decode_swa.py` around lines 417 - 427, Cache the result of init_swa_metadata so swa_indices_and_lens, init_start_pos, and init_block_table execute only once. Update init_swa_indices and init_swa_lens to reuse the cached metadata while preserving their respective contiguous outputs.
🤖 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 `@docs/models/deepseek.md`:
- Line 95: Update the entry-point table in the DeepSeek documentation to replace
the old *_attention_csa.py, *_attention_hca.py, and *_attention_swa.py
references with the renamed mode-first modules, including the applicable _cp
variants, while retaining prefill_csa.py where already correct.
In `@models/deepseek/v4-flash/decode_hca.py`:
- Around line 225-241: Update the top-k token iteration around topk_all so the
final partial HCA_TOPK_TOKEN_TILE is included rather than relying on a grid that
truncates the tail. Remove the ineffective topk_t < T guard or replace the loop
bounds with a ceil-divided token-block count, and ensure every token index below
T initializes its HCA_CMP_TOPK entries to valid compressed slots or -1
consistently with sparse_attn_hca expectations.
- Around line 13-14: Update the companion-file references in the module
docstring to use the renamed mode-first files decode_swa.py and decode_csa.py,
replacing the obsolete attention_swa.py and attention_csa_draft.py names.
- Around line 195-202: Add a module-level divisibility assertion for
HCA_WB_TOKEN_TILE, matching the guard style used in decode_sparse_attn_csa.py,
to ensure the configured T value is an exact multiple of HCA_WB_TOKEN_TILE
before the writeback loop runs. Preserve the existing writeback logic in the
hca_cache_writeback loop.
- Around line 273-276: Update attention_hca_test’s compressor-cache handling:
declare compress_state and cmp_kv as pl.InOut, and mark both corresponding
TensorSpecs with is_output=True and output comparison. Ensure
compressor_ratio128 and golden_compressor updates to these buffers are
validated.
- Around line 84-90: Replace HCA_CMP_TOPK usage with HCA_TOPK_LIMIT for the HCA
top-k bound, and add an assertion that HCA_TOPK_LIMIT is at least
get_standalone_cmp_valid(128). Do not reuse HCA_SPARSE_CMP_TOPK for this HCA
limit.
In `@models/deepseek/v4-flash/decode_sparse_attn_csa.py`:
- Line 743: Update the zip call in the loop over block_mi, block_li, and
block_oi to pass strict=True, explicitly validating that all three iterables
have equal lengths while preserving the existing iteration behavior.
- Around line 256-264: Add an assertion alongside the existing tiling-invariant
checks in the module initialization/configuration validation to require WIN ==
ATTN_K_TILE. Keep the block-validity mapping in the loops using its current
offsets, and ensure invalid configurations fail before qk_pv executes.
In `@models/deepseek/v4-flash/decode_swa.py`:
- Around line 14-15: Update the module docstring’s companion file references to
use decode_csa.py and decode_hca.py instead of the renamed
attention_csa_draft.py and attention_hca_draft.py entries, preserving the
existing ratio annotations.
In `@models/deepseek/v4-flash/prefill_csa_cp.py`:
- Line 1973: Update the sparse_bias initialization in the prefill sparse
attention path to import FP32_NEG_INF from config and use it instead of the
literal -3.0e38, keeping the existing tensor shape, dtype, and initialization
behavior unchanged.
In `@models/deepseek/v4-flash/prefill_csa.py`:
- Around line 231-237: The sparse-attention launch following
`_prefill_compressor_ratio4_with_completion` must depend on the compressor write
completion before reading `cmp_kv`. Update the
`_prefill_sparse_attn_with_block_mask` call or its launch dependency arguments
to include `compressor_completion[0]`, preserving the same CP ordering behavior
used in `prefill_csa_cp.py`.
In `@models/deepseek/v4-flash/prefill_hca_cp.py`:
- Around line 784-807: Update the HCA compression flow around
prefill_compressor_ratio128 and cp_hca_pack_compact to collect the compressor
completion IDs for each launched leaf/state write, then pass or commission those
completions through the compact payload dependency path before
cp_hca_pack_compact reads leaf_cmp_flat and scratch_state_flat. Preserve the
existing buffer slicing and packing behavior while ensuring compact packing
waits on all compressor roots, matching the CSA ratio-4 caller’s dependency
handling.
In `@models/deepseek/v4-flash/prefill_swa_cp.py`:
- Around line 306-310: Populate valid_block_mask in prefill_cp_swa using the
same mask-construction logic as _prefill_cp_sparse_stage before invoking
_prefill_sparse_attn_math. Ensure every staged attention block row is enabled
according to the SWA CP staging inputs, rather than leaving nonzero rows masked
by the zero-initialized valid_mask.
---
Outside diff comments:
In `@models/deepseek/v4-flash/prefill_indexer_compressor.py`:
- Around line 347-393: Add an explicit dependency on cache_write_tid to the
prefill_idx_c4_scale_scatter task declaration so it waits for the producer that
writes scale_scratch. Preserve the existing scale scatter logic and dependency
behavior for all other tasks.
---
Nitpick comments:
In `@models/deepseek/v4-flash/decode_sparse_attn_csa.py`:
- Around line 52-57: Remove the duplicate MAX_SEQ_LEN assignment from the CSA
compressed-slot masking block, while retaining the existing assignment near line
39. Keep INDEXER_SCORE_LEN and the remaining CSA constants using the shared
MAX_SEQ_LEN value.
- Around line 690-710: Rename the loop variable raw in both index-iteration
loops to distinct descriptive names, updating each corresponding reference while
preserving the existing index handling and masking behavior.
- Around line 542-545: Move the T_PAD padding initialization out of the for qt
token-loop and perform it once before the loop, preserving the existing
zero_half, zero_i8, and o_r_i8_pad assembly behavior for T_PAD > T. Keep
token-dependent processing inside the loop and ensure the padding rows remain
unchanged for every tile.
- Around line 316-337: Optimize the window-prefix path in the loop around qk_kv
by detecting contiguous GATHER_RUN sub-tiles, probing the first and last window
slots as in decode_sparse_attn_swa.sparse_attn_swa and
decode_sparse_attn_hca.sparse_attn_hca, and using one bulk gather when their
separation is GATHER_RUN - 1. Preserve the existing per-row gather and fallback
behavior for non-contiguous or invalid slots, without changing results for the
comparison-sparse path.
In `@models/deepseek/v4-flash/decode_swa.py`:
- Around line 417-427: Cache the result of init_swa_metadata so
swa_indices_and_lens, init_start_pos, and init_block_table execute only once.
Update init_swa_indices and init_swa_lens to reuse the cached metadata while
preserving their respective contiguous outputs.
In `@models/deepseek/v4-flash/prefill_csa_cp.py`:
- Around line 271-289: Reduce _build_metadata_tensors to return only the
metadata tables and context entries consumed by build_tensor_specs, removing the
unused boundary, seed, slot-mapping, state-mapping, and dense-prefix
calculations. Move validate_cp_indexer_capacity out of _build_metadata_tensors
and invoke it directly from build_tensor_specs, preserving the existing capacity
validation behavior independently of metadata construction.
In `@models/deepseek/v4-flash/prefill_csa.py`:
- Around line 251-259: Update the assignment receiving outputs from
prefill_indexer so the intentionally unused idx_kv_cache_out and
idx_kv_scale_out bindings use underscore-prefixed names, while preserving the
remaining outputs and call behavior.
In `@models/deepseek/v4-flash/prefill_exchange_cp.py`:
- Around line 478-623: Refactor the SWA path to reuse _prefill_cp_sparse_stage
instead of maintaining the duplicate _cp_swa_stage_sources implementation in
prefill_swa_cp.py. Supply IDX_TOPK-wide cmp_indices initialized to -1 and the
required compressed-stage buffers so the compressed branch remains inert, while
preserving the existing SWA outputs and overlay behavior.
In `@models/deepseek/v4-flash/prefill_hca_cp.py`:
- Around line 647-662: In cp_hca_tail_assemble, remove the redundant if/else
around source and assign it directly using part * MAX_SEGMENT_TILES * TAIL_ROWS
+ tail_offset. Apply the same simplification to the corresponding tail assembly
logic in prefill_swa_cp.py, preserving the existing contiguous tile layout and
indexing behavior.
- Around line 1313-1314: Derive the shared segment geometry once before the
metadata construction in the fixture setup, then pass that geometry into both
build_hca_metadata and _build_raw_attention_metadata instead of recomputing
prefix, span, lengths, starts, owners, and query_positions independently. Update
both builder signatures and their callers, and ensure device_values uses fields
from this single consistent geometry source.
In `@models/deepseek/v4-flash/prefill_indexer.py`:
- Around line 379-482: Extract the duplicated query-preparation pipeline from
`_prefill_indexer_cp_score_topk` and `prefill_indexer` into one `@pl.jit.inline`
helper that returns `qr_hadamard_i8` and `qr_hadamard_scale_dq`. Move the shared
projection/dequantization, RoPE preparation and rotation, Hadamard
multiplication, and per-row INT8 quantization into the helper, using a
`name_prefix` parameter to generate task-specific `name_hint` values, then
replace each duplicated block with the helper call.
In `@models/deepseek/v4-flash/prefill_swa_cp.py`:
- Around line 502-510: Rename the local scalar `tail_start` in `prefill_cp_swa`
to `tail_start0`, and update its use when computing `tail_offset`; preserve the
existing tail-copy behavior while avoiding shadowing the module-level
`tail_start` helper.
- Around line 642-644: Update build_tensor_specs to reject any cp_size that
differs from the import-time CP_SIZE before building metadata, set
golden_prefill_cp_swa._ctx within the builder, and return only specs. Adjust the
__main__ block to use the builder’s new return contract without externally
installing _ctx.
In `@models/deepseek/v4-flash/prefill_zigzag_cp.py`:
- Around line 33-43: Export and reuse the existing _parse_static_int helper from
prefill_zigzag_cp.py in prefill_swa_cp.py, removing the duplicate implementation
and deriving CP_SIZE through the shared parser (or importing CP_SIZE directly).
Ensure both modules use one consistent --cp argument parsing rule.
🪄 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 Plus
Run ID: 3f6cfc98-1646-4186-8e55-061d272ea038
📒 Files selected for processing (27)
docs/debug-and-tune/debugging.mddocs/debug-and-tune/performance-tuning.mddocs/models/deepseek.mddocs/run-and-validate/compile-runtime-workflow.mdmodels/deepseek/v4-flash/decode_csa.pymodels/deepseek/v4-flash/decode_fwd.pymodels/deepseek/v4-flash/decode_hca.pymodels/deepseek/v4-flash/decode_indexer.pymodels/deepseek/v4-flash/decode_layer.pymodels/deepseek/v4-flash/decode_mtp.pymodels/deepseek/v4-flash/decode_sparse_attn_csa.pymodels/deepseek/v4-flash/decode_swa.pymodels/deepseek/v4-flash/prefill_compressor_ratio4.pymodels/deepseek/v4-flash/prefill_csa.pymodels/deepseek/v4-flash/prefill_csa_cp.pymodels/deepseek/v4-flash/prefill_exchange_cp.pymodels/deepseek/v4-flash/prefill_fwd.pymodels/deepseek/v4-flash/prefill_hca.pymodels/deepseek/v4-flash/prefill_hca_cp.pymodels/deepseek/v4-flash/prefill_indexer.pymodels/deepseek/v4-flash/prefill_indexer_compressor.pymodels/deepseek/v4-flash/prefill_layer.pymodels/deepseek/v4-flash/prefill_mtp.pymodels/deepseek/v4-flash/prefill_sparse_attn.pymodels/deepseek/v4-flash/prefill_swa.pymodels/deepseek/v4-flash/prefill_swa_cp.pymodels/deepseek/v4-flash/prefill_zigzag_cp.py
|
|
||
| ```bash | ||
| python models/deepseek/v4-flash/prefill_attention_csa.py -p a2a3sim | ||
| python models/deepseek/v4-flash/prefill_csa.py -p a2a3sim |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the entry-point table with the renamed attention modules.
This command now uses prefill_csa.py, but the same document still lists *_attention_csa.py, *_attention_hca.py, and *_attention_swa.py. Update the table to use the mode-first module names and include _cp variants where applicable.
As per coding guidelines, keep public documentation and examples aligned when behavior changes.
🤖 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 `@docs/models/deepseek.md` at line 95, Update the entry-point table in the
DeepSeek documentation to replace the old *_attention_csa.py,
*_attention_hca.py, and *_attention_swa.py references with the renamed
mode-first modules, including the applicable _cp variants, while retaining
prefill_csa.py where already correct.
Source: Coding guidelines
| compressor_completion = pl.array.create(1, pl.TASK_ID) | ||
| _prefill_compressor_ratio4_with_completion( | ||
| x_normed, compress_state, compress_state_block_table, | ||
| cmp_wkv, cmp_wgate, cmp_ape, cmp_norm_w, | ||
| freqs_cos, freqs_sin, cmp_kv, | ||
| position_ids, num_tokens, cmp_slot_mapping, state_slot_mapping, | ||
| cmp_wkv, cmp_wgate, cmp_ape, | ||
| cmp_norm_w, freqs_cos, freqs_sin, | ||
| cmp_kv, position_ids, num_tokens, | ||
| cmp_slot_mapping, state_slot_mapping, compressor_completion, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n -C 10 \
'_prefill_compressor_ratio4_with_completion|compressor_completion|_prefill_sparse_attn_with_block_mask' \
models/deepseek/v4-flash/prefill_csa.py \
models/deepseek/v4-flash/prefill_compressor_ratio4.py
rg -n -C 8 \
'main_completion|inner_completion|deps=.*completion' \
models/deepseek/v4-flash/prefill_csa_cp.pyRepository: hw-native-sys/pypto-lib
Length of output: 15138
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== prefill_csa.py relevant sections =="
sed -n '200,335p' models/deepseek/v4-flash/prefill_csa.py
echo
echo "== prefill_compressor_ratio4.py signature/body tail =="
sed -n '55,220p' models/deepseek/v4-flash/prefill_compressor_ratio4.py
sed -n '380,416p' models/deepseek/v4-flash/prefill_compressor_ratio4.py
echo
echo "== prefill_sparse_attn files =="
fd -a 'prefill_sparse_attn.py$' .
for f in $(fd 'prefill_sparse_attn.py$' .); do
echo "--- $f outline/load/sparse_attn ---"
rg -n "def _prefill_sparse_attn_with_block_mask|prefill_csa.py|at[level=|deps=" "$f" -C 4 || true
done
echo
echo "== task dependency / pl.at usage in deepseek v4-flash =="
rg -n "pl\.at\(|deps=.*completion|completion\[0\]|main_completion|inner_completion" models/deepseek/v4-flash -C 3
echo
echo "== pypto imports and TASK_ID creation semantics in repo =="
rg -n "TASK_ID|create\\(1, pl\\.TASK_ID\\)|array\\.create|pl\\.array" --iglob '*.md' --iglob '*.py' | head -n 200Repository: hw-native-sys/pypto-lib
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== prefill_sparse_attn.py signature/body outline =="
awk '/def _prefill_sparse_attn_with_block_mask/,/^def / {print NR": "$0}' models/deepseek/v4-flash/prefill_sparse_attn.py | head -n 220
echo
echo "== exact sparse attn call in prefill_csa.py =="
sed -n '300,332p' models/deepseek/v4-flash/prefill_csa.py
echo
echo "== cmp_kv producers in prefill_compressor_ratio4.py =="
rg -n "cmp_kv|cmp_kv_flat|cmp.*write|write_dst|pl\.write|cmp_cache" models/deepseek/v4-flash/prefill_compressor_ratio4.py -C 4
echo
echo "== sparse_attn reads/writes related to cmp_kv in prefill_sparse_attn.py =="
rg -n "cmp_kv|sparse_kv|write|pl\.read|pl\.assemble|tile\.get_block_idx|spmd" models/deepseek/v4-flash/prefill_sparse_attn.py -C 5 | head -n 260
echo
echo "== deterministic task dependency check: does sparse launch include compressor completion deps? =="
python3 - <<'PY'
from pathlib import Path
import re
csa = Path("models/deepseek/v4-flash/prefill_csa.py").read_text()
attn = Path("models/deepseek/v4-flash/prefill_sparse_attn.py").read_text()
def text_between(start, end):
i = csa.find(start)
j = csa.find(end, i + 1)
return csa[i:j] if i >= 0 and j > 0 else ""
call_block = text_between("compressor_completion = pl.array.create", "attn_out = pl.create_tensor").split("_prefill_sparse_attn_with_block_mask")[0] + "attn_out = pl.create_tensor"
print("has compressor_completion creation:", "compressor_completion = pl.array.create(1, pl.TASK_ID)" in csa)
print("has compressor_completion passed to compressor:", "_prefill_compressor_ratio4_with_completion(" in call_block and "compressor_completion," in call_block)
print("has compressor_completion read/consumed by sparse block or later:", bool(re.search(r'compressor_completion\s*\[\s*0\s*\]', csa)))
call = re.search(r'_prefill_sparse_attn_with_block_mask\s*\((.*?)\n\s*\)', csa, re.S)
call_text = call.group(1) if call else ""
print("sparse call lines:", csa[call.start():call.end()+1] if call else "NOT FOUND")
print("sparse call contains 'deps':", "deps" in call_text)
print("sparse call contains any completion id:", bool(re.search(r'\.completion|completion|_tid\s*\)', call_text)))
print("sparse file contains deps in at/with context:", bool(re.search(r'\bdeps\s*=', attn)))
PYRepository: hw-native-sys/pypto-lib
Length of output: 27950
🌐 Web query:
PyPTO pl.at deps pl.system.task_dummy Tensor dependency completion documentation
💡 Result:
In the PyPTO (Parallel Tensor/Tile Operation) framework, task dependencies are primarily managed automatically by the runtime (specifically the TensorMap system) based on tensor read/write patterns [1]. However, developers can manually influence task scheduling and dependency edges when automatic inference is insufficient or non-deterministic [2][3]. Regarding the specific components mentioned: 1. pl.at: This is a scope-based construct used within PyPTO to define execution contexts (e.g., specifying pl.Level.CORE_GROUP) and hint at task boundaries or kernel optimization strategies [4][5][6]. 2. pl.system.task_dummy: This refers to a placeholder "dummy" task that was historically used in PyPTO orchestration to force a bridge or synchronization point between two operations (e.g., bridging an attention layer to an output projection) when the automatic dependency graph could not detect the necessary ordering [3]. Current best practices in PyPTO development involve replacing these dummy tasks with explicit dependency management [3]. 3. Tensor Dependency Completion and Manual Overrides: When the automatic dependency tracker fails to establish a required edge (such as in Write-After-Read/WAR scenarios where the runtime fails to infer a serialization dependency on a shared buffer), developers use explicit dependency injection [2]. The framework exposes pl.submit(kernel, *args, deps=[...]) for this purpose [2]. By passing a list of deps (containing producer TaskIDs), the developer explicitly mandates that a task cannot start until its listed predecessors have completed [2][3]. This approach is preferred over using dummy tasks or "no-op" hacks (like forced self-copies) because it provides precise, overhead-free control over the dependency graph [2][3]. In summary, while PyPTO aims for automated dependency tracking through tensor dataflow analysis, developers should utilize explicit deps arguments in pl.submit to handle complex scheduling requirements or to resolve non-deterministic race conditions that the runtime cannot automatically detect [2][1].
Citations:
- 1: https://github.com/xwhu/pypto-lib/blob/main/pto2_rt.md
- 2: [Bug] decode_attention_swa/hca need a no-op self-copy hack to force a WAR edge auto-dep misses — replace with pl.submit(deps=) manual dependency #481
- 3: Refactor: drop dummy attn_fence and tighten Qwen3-14B LM head #489
- 4: https://github.com/hw-native-sys/pypto-lib
- 5: e84fb6f
- 6: https://github.com/hw-native-sys/pypto-lib/blob/main/README.md
Wait for the compressor completion before starting sparse attention.
_prefill_compressor_ratio4_with_completion writes cmp_kv back through cmp_kv_flat and exposes the write completion in compressor_completion, but _prefill_sparse_attn_with_block_mask reads cmp_kv without that dependency. Add compressor_completion[0] to the sparse-attention launch/dependency so the CP ordering in models/deepseek/v4-flash/prefill_csa_cp.py also applies here.
🤖 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/prefill_csa.py` around lines 231 - 237, The
sparse-attention launch following `_prefill_compressor_ratio4_with_completion`
must depend on the compressor write completion before reading `cmp_kv`. Update
the `_prefill_sparse_attn_with_block_mask` call or its launch dependency
arguments to include `compressor_completion[0]`, preserving the same CP ordering
behavior used in `prefill_csa_cp.py`.
| for part in pl.range(LOCAL_PARTS): | ||
| state_base = part * HCA_STATE_PHYSICAL_BLOCKS | ||
| state_end = state_base + HCA_STATE_PHYSICAL_BLOCKS | ||
| state_part = scratch_state[state_base:state_end, 0:HCA_STATE_BLOCK_SIZE, 0:COMPRESS_STATE_DIM] | ||
| for leaf in pl.range(MAX_COMPRESS_LEAVES): | ||
| leaf_index = part * MAX_COMPRESS_LEAVES + leaf | ||
| token0 = leaf_index * TAIL_ROWS | ||
| cmp_block0 = leaf_index * LEAF_CMP_BLOCKS | ||
| x_leaf = effective_x[token0 : token0 + TAIL_ROWS, 0:D] | ||
| position_leaf = leaf_positions[token0 : token0 + TAIL_ROWS] | ||
| cmp_slots_leaf = leaf_cmp_slots[token0 : token0 + TAIL_ROWS] | ||
| state_slots_leaf = leaf_state_slots[token0 : token0 + TAIL_ROWS] | ||
| cmp_leaf = leaf_cmp[cmp_block0 : cmp_block0 + LEAF_CMP_BLOCKS, 0:BLOCK_SIZE, 0:1, 0:HEAD_DIM] | ||
| active = pl.read(leaf_num_tokens, [part, leaf]) | ||
| prefill_compressor_ratio128( | ||
| x_leaf, state_part, compress_state_block_table, | ||
| cmp_wkv, cmp_wgate, cmp_ape, cmp_norm_w, | ||
| freqs_cos, freqs_sin, | ||
| cmp_leaf, position_leaf, active, | ||
| cmp_slots_leaf, state_slots_leaf, | ||
| ) | ||
| leaf_cmp[cmp_block0 : cmp_block0 + LEAF_CMP_BLOCKS, 0:BLOCK_SIZE, 0:1, 0:HEAD_DIM] = cmp_leaf | ||
| state_end = state_base + HCA_STATE_PHYSICAL_BLOCKS | ||
| scratch_state[state_base:state_end, 0:HCA_STATE_BLOCK_SIZE, 0:COMPRESS_STATE_DIM] = state_part |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare prefill_compressor_ratio128 and the ratio-4 completion-aware variant.
set -euo pipefail
fd -t f 'prefill_compressor_ratio128.py' models --exec ast-grep outline {} --items all \;
fd -t f 'prefill_compressor_ratio128.py' models --exec rg -n 'def prefill_compressor_ratio128|completion|task_dummy|return ' {} \;
rg -nP -C3 'prefill_compressor_ratio128\(' --type=py -g 'models/**'Repository: hw-native-sys/pypto-lib
Length of output: 5807
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== file sizes =="
wc -l models/deepseek/v4-flash/prefill_hca_cp.py models/deepseek/v4-flash/prefill_csa_cp.py models/deepseek/v4-flash/prefill_compressor_ratio128.py 2>/dev/null || true
echo "== hca cp outline around relevant symbols =="
ast-grep outline models/deepseek/v4-flash/prefill_hca_cp.py --view expanded | sed -n '1,220p'
echo "== hca cp lines 740-920 =="
sed -n '740,920p' models/deepseek/v4-flash/prefill_hca_cp.py | nl -ba -v740
echo "== hca pack references =="
rg -n -C4 'cp_hca_pack_compact|leaf_cmp_flat|scratch_state_flat|task_dummy|completion|return .*cmp_kv|return .*compress_state|prefill_compressor_ratio128' models/deepseek/v4-flash/prefill_hca_cp.py
echo "== ratio128 definition lines 72-284 =="
sed -n '72,284p' models/deepseek/v4-flash/prefill_compressor_ratio128.py | nl -ba -v72Repository: hw-native-sys/pypto-lib
Length of output: 2678
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== csa cp lines around completion usage =="
sed -n '1050,1115p' models/deepseek/v4-flash/prefill_csa_cp.py | nl -ba -v1050
echo "== csa cp compressor call context if any =="
rg -n -C5 'prefill_compressor_ratio128|leaf_cmp_flat|scratch_state_flat|cp_hca_pack_compact|task_dummy|completion' models/deepseek/v4-flash/prefill_csa_cp.py || trueRepository: hw-native-sys/pypto-lib
Length of output: 244
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== hca cp lines 740-920 =="
python3 - <<'PY'
from pathlib import Path
p=Path('models/deepseek/v4-flash/prefill_hca_cp.py')
text=p.read_text().splitlines()
for i in range(739, min(919, len(text))):
print(f"{i+1}: {text[i]}")
PY
echo "== hca cp pack and compressor hits =="
rg -n -C4 'cp_hca_pack_compact|leaf_cmp_flat|scratch_state_flat|task_dummy|completion|prefill_compressor_ratio128' models/deepseek/v4-flash/prefill_hca_cp.py
echo "== csa cp lines 1050-1125 =="
python3 - <<'PY'
from pathlib import Path
p=Path('models/deepseek/v4-flash/prefill_csa_cp.py')
text=p.read_text().splitlines()
for i in range(1049, min(1124, len(text))):
print(f"{i+1}: {text[i]}")
PYRepository: hw-native-sys/pypto-lib
Length of output: 20631
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== ratio128 definition lines 72-284 =="
python3 - <<'PY'
from pathlib import Path
p=Path('models/deepseek/v4-flash/prefill_compressor_ratio128.py')
text=p.read_text().splitlines()
for i in range(71, min(283, len(text))):
print(f"{i+1}: {text[i]}")
PY
echo "== ratio4 completion-aware definitions =="
rg -n -C4 'def _prefill_compressor_ratio4_with_completion|prefill_compressor_ratio4_with_completion|completion|task_dummy|return .*main_cache_written|return .*main_state_written' models/deepseek/v4-flash models/deepseek/v4-pro --type=py || true
echo "== return/completion references in ratio128 files =="
rg -n -C2 'return cmp_kv, compress_state|task_dummy|completion|array\.create|with pl\.at\(.*deps|return.*written' models/deepseek/v4-flash/prefill_compressor_ratio128.py models/deepseek/v4-pro/prefill_compressor_ratio128.py || true
echo "== all ratio128 call sites =="
rg -n -C3 'prefill_compressor_ratio128\(' models --type=py || trueRepository: hw-native-sys/pypto-lib
Length of output: 50379
Gate compact packing behind the compressor completions.
prefill_compressor_ratio128 writes the caller-owned cmp_kv and compress_state buffers through cmp_leaf and state_part, but it does not return a completion ID. cp_hca_pack_compact then reads leaf_cmp_flat and scratch_state_flat without a dependency on those writes, unlike the CSA ratio-4 caller. Add a completion-array dependency for the HCA compressor and capture/commission the compressor root returns through the compact payload path.
🤖 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/prefill_hca_cp.py` around lines 784 - 807, Update
the HCA compression flow around prefill_compressor_ratio128 and
cp_hca_pack_compact to collect the compressor completion IDs for each launched
leaf/state write, then pass or commission those completions through the compact
payload dependency path before cp_hca_pack_compact reads leaf_cmp_flat and
scratch_state_flat. Preserve the existing buffer slicing and packing behavior
while ensuring compact packing waits on all compressor roots, matching the CSA
ratio-4 caller’s dependency handling.
| sparse_kv: pl.Tensor[[LOCAL_SPARSE_ROWS, HEAD_DIM], pl.BF16], | ||
| sparse_bias: pl.Tensor[[NUM_LOCAL_TILES * TAIL_ROWS, PREFILL_SPARSE_PAD], pl.FP32], | ||
| valid_block_mask: pl.Tensor[[NUM_LOCAL_TILES * TAIL_ROWS, VALID_BLOCK_MASK_COLS], pl.INT32], | ||
| overlay_active_lengths: pl.Tensor[[NUM_LOCAL_TILES, OVERLAY_SOURCES], pl.INT32], | ||
| ): |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Resolve WIN, PREFILL_ATTN_TILE, PREFILL_ATTN_BLOCKS, and VALID_BLOCK_MASK_COLS.
set -euo pipefail
fd -t f 'prefill_sparse_attn.py' models/deepseek --exec rg -n 'PREFILL_ATTN_TILE\s*=|PREFILL_ATTN_BLOCKS\s*=|PREFILL_SPARSE_PAD\s*=|SPARSE_BIAS_COLS\s*=|VALID_BLOCK_MASK_COLS\s*=|^WIN\s*=' {} \;
fd -t f 'config.py' models/deepseek --exec rg -n 'sliding_window' {} \;
# Confirm no other site writes valid_mask in the SWA CP path.
fd -t f 'prefill_swa_cp.py' models/deepseek --exec rg -n 'valid_mask|valid_block_mask' {} \;Repository: hw-native-sys/pypto-lib
Length of output: 1381
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== files =="
fd -t f 'prefill_swa_cp.py|prefill_sparse_attn.py|prefill_exchange_cp.py|config.py' models/deepseek
echo "== prefill_swa_cp.py outline =="
ast-grep outline models/deepseek/prefill_swa_cp.py --view expanded || true
wc -l models/deepseek/prefill_swa_cp.py
echo "== relevant prefill_swa_cp.py sections =="
sed -n '270,580p' models/deepseek/prefill_swa_cp.py | cat -n
echo "== relevant prefill_sparse_attn.py sections =="
sed -n '35,140p' models/deepseek/prefill_sparse_attn.py | cat -n
sed -n '180,320p' models/deepseek/prefill_sparse_attn.py | cat -n
sed -n '320,460p' models/deepseek/prefill_sparse_attn.py | cat -n
echo "== prefill_exchange_cp.py around mask build =="
sed -n '560,645p' models/deepseek/prefill_exchange_cp.py | cat -n
echo "== sliding window config occurrences =="
sed -n '1,35p' models/deepseek/config.py | cat -nRepository: hw-native-sys/pypto-lib
Length of output: 618
🏁 Script executed:
#!/bin/bash
set -euo pipefail
FILE="models/deepseek/v4-flash/prefill_swa_cp.py"
OTHER="models/deepseek/v4-flash/prefill_sparse_attn.py"
EX="models/deepseek/v4-flash/prefill_exchange_cp.py"
CONFIG="models/deepseek/v4-flash/config.py"
echo "== file sizes =="
wc -l "$FILE" "$OTHER" "$EX" "$CONFIG"
echo "== relevant prefill_swa_cp.py sections =="
sed -n '250,590p' "$FILE" | cat -n
echo "== relevant prefill_sparse_attn.py sections =="
sed -n '1,180p' "$OTHER" | cat -n
sed -n '180,460p' "$OTHER" | cat -n
echo "== prefill_exchange_cp.py around mask build =="
sed -n '560,635p' "$EX" | cat -n
echo "== config window and helper sections =="
sed -n '1,180p' "$CONFIG" | cat -n
echo "== deterministic config derivation for a2a3/a2a3sim/a5/a5sim =="
python3 - <<'PY'
from pathlib import Path
import re, ast
config_path = Path("models/deepseek/v4-flash/config.py")
text = config_path.read_text()
config = {}
exec(text, config)
for name in ["a2a3", "a2a3sim", "a5", "a5sim"]:
m = config.get(name)
print(name, m.sliding_window if m else "*", "PREFILL_ATTN_TILE", config.get("PREFILL_ATTN_TILE"), "PREFILL_SPARSE_TOPK", config.get("PREFILL_SPARSE_TOPK"))
sw = m.sliding_window if m else None
attn_tile = config.get("PREFILL_ATTN_TILE")
topk = config.get("PREFILL_SPARSE_TOPK")
if sw is None or attn_tile is None or topk is None:
print(name, "undetermined")
continue
blocks = (topk + attn_tile - 1) // attn_tile
print(name, "sw_blocks", blocks if sw > attn_tile else 1,
"topk_tiles", (topk + attn_tile - 1) // attn_tile)
PYRepository: hw-native-sys/pypto-lib
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
CONFIG="models/deepseek/v4-flash/config.py"
echo "== model configs with keys =="
python3 - <<'PY' 2>&1
from pathlib import Path
import ast
path = Path("models/deepseek/v4-flash/config.py")
tree = ast.parse(path.read_text())
classes = {n: [] for n in ("Flash", "FlashV2")}
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name in classes:
for item in node.body:
if isinstance(item, ast.Assign):
for target in item.targets:
if isinstance(target, ast.Name):
classes[node.name].append(target.id)
for name, keys in classes.items():
print(name, " ".join(key for key in keys if key.lower().startswith(("sliding_", "index_", "flash"))))
# Print relevant class bodies compactly.
for name in ("Flash", "FlashV2"):
print(f"\n== {name} relevant attributes ==")
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name == name:
for item in node.body:
if isinstance(item, ast.Assign):
ids = " ".join(t.id if isinstance(t, ast.Name) else ast.unparse(t) for t in item.targets)
val = ast.unparse(item.value)
if any(k in ids for k in ("sliding_window", "index_topk", "o_groups", "head_dim", "name")):
print(f"{ids} = {val}")
PY
echo "== deterministic window/topk configuration =="
python3 - <<'PY'
from pathlib import Path
import ast
src = Path("models/deepseek/v4-flash/config.py").read_text().splitlines()
env = {}
lines = []
for i, line in enumerate(src, start=1):
if "VALID_BLOCK_MASK_COLS" in line or "PREFILL_ATTN_BLOCKS" in line or "PREFILL_SPARSE_PAD" in line:
print(f"{i}: {line}")
PYRepository: hw-native-sys/pypto-lib
Length of output: 321
🏁 Script executed:
#!/bin/bash
set -euo pipefail
CONFIG="models/deepseek/v4-flash/config.py"
SOURCER="models/deepseek/v4-flash/prefill_swa_cp.py"
MATH="models/deepseek/v4-flash/prefill_sparse_attn.py"
echo "== top-level names =="
grep -nE '^[A-Z][A-Za-z0-9_]* = ' "$CONFIG" "$SOURCER" "$MATH" | rg 'Flash|sliding_window|sliding|index|top|local|local' || true
echo "== all valid_block_mask writes in deepseek/v4-flash Python files =="
rg -n 'valid_block_mask|valid_mask|prefill_cp_build_valid_mask|prefill_cp_sparse_stage' models/deepseek/v4-flash -g '*.py' || true
echo "== Flash config text =="
cat "$CONFIG" | sed -n '1,230p'
echo "== relevant _cp_swa_stage_sources references around write candidates =="
python3 - <<'PY'
from pathlib import Path
p = Path("models/deepseek/v4-flash/prefill_swa_cp.py")
text = p.read_text().splitlines()
needles = [
"pl.write(mask",
"valid_block_mask",
"sparse_kv",
]
for start, end in [(40, 130), (435, 530)]:
print(f"\n---- {p}:{start}-{end} ----")
for i in range(start, min(end, len(text))+1):
if any(n in text[i-1] for n in needles):
print(f"{i:4}: {text[i-1]}")
PYRepository: hw-native-sys/pypto-lib
Length of output: 14185
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== sparse topk and blocks in v4-flash config via Python eval =="
python3 - <<'PY'
from pathlib import Path
import ast
src = Path("models/deepseek/v4-flash/config.py").read_text()
top_level = "\n".join(line for line in src.splitlines() if not line.strip().startswith("#"))
ns = {}
exec(compile(top_level, "config.py", "exec"), ns)
flash = ns["FLASH"]
topk_window_entries = flash.sliding_window + flash.index_topk
topk_window_entries = min(topk_window_entries, len(flash.compress_ratios))
attn_tile = 128
topk = slide_add_idx_topk = flash.sliding_window + flash.index_topk
sparse_topk = min(slide_add_idx_topk, min(flash.sliding_window, flash.max_position_embeddings) + min(1, min(flash.index_topk, flash.sliding_window + flash.sliding_window // 2)))
blocks = (sparse_topk + attn_tile - 1) // attn_tile
print("sliding_window", flash.sliding_window)
print("index_topk", flash.index_topk)
print("topk_window_entries", topk_window_entries)
print("sparse_topk", sparse_topk)
print("blocks", blocks)
print("PREFILL_SPARSE_PAD", blocks * attn_tile)
PYRepository: hw-native-sys/pypto-lib
Length of output: 330
Populate valid_block_mask in the SWA CP stage.
_cp_swa_stage_sources stages more than one attention block for the default config, but prefill_cp_swa creates valid_mask with init_value=0 and passes it only to _prefill_sparse_attn_math. In that backend, block 0 is force-enabled, but the remaining sparse block rows stay masked. Add the same mask build used by _prefill_cp_sparse_stage before calling the math backend.
🤖 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/prefill_swa_cp.py` around lines 306 - 310, Populate
valid_block_mask in prefill_cp_swa using the same mask-construction logic as
_prefill_cp_sparse_stage before invoking _prefill_sparse_attn_math. Ensure every
staged attention block row is enabled according to the SWA CP staging inputs,
rather than leaving nonzero rows masked by the zero-initialized valid_mask.
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 13
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/prefill_indexer_compressor.py (1)
347-393: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAdd an explicit dependency from
prefill_idx_c4_scale_scattertocache_write_tid.
scale_scratchis written inside theprefill_idx_c4_cache_writeSPMD task at line 368. Theprefill_idx_c4_scale_scatterblock at line 382 readsscale_scratchbut declares nodeps. The producer task idcache_write_tidis already in scope and is only used later at line 425 for the completion signal.If the scheduler does not derive this read-after-write edge automatically, the scatter can publish stale or uninitialized dequantization scales into
idx_kv_scale. That corrupts the C8 indexer cache while the INT8 rows stay correct, which makes the defect hard to trace.🔒️ Proposed fix to order the scale scatter after the cache write
- with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_idx_c4_scale_scatter") as scale_scatter_tid: + with pl.at( + level=pl.Level.CORE_GROUP, + name_hint="prefill_idx_c4_scale_scatter", + deps=[cache_write_tid], + ) as scale_scatter_tid:🤖 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/prefill_indexer_compressor.py` around lines 347 - 393, Add an explicit dependency on cache_write_tid to the prefill_idx_c4_scale_scatter task declaration so it waits for the producer that writes scale_scratch. Preserve the existing scale scatter logic and dependency behavior for all other tasks.
🧹 Nitpick comments (14)
models/deepseek/v4-flash/prefill_csa.py (1)
251-259: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMark the unused indexer outputs as intentional.
Ruff RUF059 reports
idx_kv_cache_outandidx_kv_scale_outas unused. Ifprefill_indexerupdatesidx_kv_cacheandidx_kv_scalein place, use underscore-prefixed bindings.Suggested binding change
- idx_kv_cache_out, idx_kv_scale_out, idx_score_unused, cmp_topk_indices = prefill_indexer( + _idx_kv_cache_out, _idx_kv_scale_out, idx_score_unused, cmp_topk_indices = prefill_indexer(🤖 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/prefill_csa.py` around lines 251 - 259, Update the assignment receiving outputs from prefill_indexer so the intentionally unused idx_kv_cache_out and idx_kv_scale_out bindings use underscore-prefixed names, while preserving the remaining outputs and call behavior.Source: Linters/SAST tools
models/deepseek/v4-flash/prefill_swa_cp.py (2)
502-510: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueThe loop variable
tail_startshadows the module-leveltail_startfunction.Line 502 binds
tail_startto a scalar insideprefill_cp_swa. Line 142 definestail_start(seg_start, seg_len)at module scope. The function is not called later in this kernel, so the shadowing is currently harmless.Rename the local to
tail_start0, matchingtail_offset0inprefill_hca_cp.pyline 646. That prevents a confusing failure if this kernel later needs the helper.🤖 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/prefill_swa_cp.py` around lines 502 - 510, Rename the local scalar `tail_start` in `prefill_cp_swa` to `tail_start0`, and update its use when computing `tail_offset`; preserve the existing tail-copy behavior while avoiding shadowing the module-level `tail_start` helper.
642-644: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAlign
build_tensor_specswith the sibling CP modules and add thecp_sizeguard.This builder returns
(specs, ctx)at line 704 and the__main__block installsgolden_prefill_cp_swa._ctxexternally at line 854.prefill_csa_cp.pyline 1007 andprefill_hca_cp.pyline 1546 instead set_ctxinside the builder and return only the spec list. A shared harness that callsbuild_tensor_specsuniformly across these modules will break on this one.This builder also omits the
cp_sizeguard.prefill_hca_cp.pylines 1309-1312 andprefill_csa_cp.pylines 639-642 raise when the runtimecp_sizediffers from the import-timeCP_SIZE. Here,build_metadata(args.cp)would build metadata for one CP size while the compiled kernel shapes use another.Set
_ctxinsidebuild_tensor_specs, return onlyspecs, and raise whencp_size != CP_SIZE.♻️ Proposed change
def build_tensor_specs(cp_size: int = CP_SIZE): + if cp_size != CP_SIZE: + raise ValueError( + f"runtime cp_size={cp_size} does not match static CP_SIZE={CP_SIZE}" + ) meta, ctx = build_metadata(cp_size)specs.append(TensorSpec("x_out", list(x.shape), torch.float32, is_output=True)) - return specs, ctx + golden_prefill_cp_swa._ctx = ctx + return specs- specs, ctx = build_tensor_specs(args.cp) - golden_prefill_cp_swa._ctx = ctx result = run_jit( fn=prefill_cp_swa_test, - specs=specs, + specs=build_tensor_specs(args.cp), golden_fn=golden_prefill_cp_swa,Also applies to: 836-854
🤖 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/prefill_swa_cp.py` around lines 642 - 644, Update build_tensor_specs to reject any cp_size that differs from the import-time CP_SIZE before building metadata, set golden_prefill_cp_swa._ctx within the builder, and return only specs. Adjust the __main__ block to use the builder’s new return contract without externally installing _ctx.models/deepseek/v4-flash/prefill_hca_cp.py (2)
647-662: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCollapse the no-op branch in
cp_hca_tail_assemble.Both arms of the
if tail_offset < TAIL_ROWSbranch compute the samesourceexpression. The branch has no effect.The layout is already correct: tiles of one part are contiguous, so
part * MAX_SEGMENT_TILES * TAIL_ROWS + tail_offsetaddresses tile 0 fortail_offset < TAIL_ROWSand tile 1 above it. The branch reads as an unfinished edit and invites a wrong "fix" later.The equivalent code in
prefill_swa_cp.pylines 506-509 has the same redundancy in a different form (TAIL_ROWS + tail_offset - TAIL_ROWS). Simplify both.♻️ Proposed simplification
for row in pl.range(TAIL_ROWS): tail_offset = tail_offset0 + row if tail_offset < total: - if tail_offset < TAIL_ROWS: - source = ( - part * MAX_SEGMENT_TILES * TAIL_ROWS - + tail_offset - ) - else: - source = ( - part * MAX_SEGMENT_TILES * TAIL_ROWS - + tail_offset - ) + # Tiles of one part are contiguous, so the offset indexes directly. + source = part * MAX_SEGMENT_TILES * TAIL_ROWS + tail_offset destination = part * TAIL_ROWS + row🤖 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/prefill_hca_cp.py` around lines 647 - 662, In cp_hca_tail_assemble, remove the redundant if/else around source and assign it directly using part * MAX_SEGMENT_TILES * TAIL_ROWS + tail_offset. Apply the same simplification to the corresponding tail assembly logic in prefill_swa_cp.py, preserving the existing contiguous tile layout and indexing behavior.
1313-1314: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffTwo metadata builders produce overlapping fixture data.
build_hca_metadataat line 322 and_build_raw_attention_metadataat line 169 both recomputeprefix,span,lengths,starts,owners, andquery_positionsfrom the same inputs.device_valuesat lines 1501-1523 then mixes the two sources:query_positionscomes fromraw_metadata, whilesegment_starts_t,cmp_indices, andsegment_tail_positionscome frommetadata.The two
query_positionstensors also differ in their inactive fill:_build_raw_attention_metadatauses0at line 177 andbuild_hca_metadatauses-1at line 335. The kernel only reads active rows, so today this is benign, but the split makes the fixture contract hard to audit.Derive the shared segment geometry once and pass it to both builders.
🤖 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/prefill_hca_cp.py` around lines 1313 - 1314, Derive the shared segment geometry once before the metadata construction in the fixture setup, then pass that geometry into both build_hca_metadata and _build_raw_attention_metadata instead of recomputing prefix, span, lengths, starts, owners, and query_positions independently. Update both builder signatures and their callers, and ensure device_values uses fields from this single consistent geometry source.models/deepseek/v4-flash/prefill_zigzag_cp.py (1)
33-43: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExport
_parse_static_intinstead of copying it intoprefill_swa_cp.py.
prefill_swa_cp.pylines 75-82 define an identical_parse_static_intand then derive its ownCP_SIZEfromCP_DEFAULT, which it already imports from this module. Two independent argv scans can disagree if the parsing rule ever changes.Import this helper in
prefill_swa_cp.py, or importCP_SIZEdirectly since both modules derive the same value from the same flag.🤖 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/prefill_zigzag_cp.py` around lines 33 - 43, Export and reuse the existing _parse_static_int helper from prefill_zigzag_cp.py in prefill_swa_cp.py, removing the duplicate implementation and deriving CP_SIZE through the shared parser (or importing CP_SIZE directly). Ensure both modules use one consistent --cp argument parsing rule.models/deepseek/v4-flash/prefill_csa_cp.py (1)
271-289: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift
_build_metadata_tensorscomputes a large unused result, and it hides the only capacity check.
build_tensor_specsat line 643 unpacksmetadata, ctxand then reads only four entries:compress_state_block_table,inner_compress_state_block_table,cmp_block_table, andidx_block_table. All four come from_build_block_tables. Ruff also reports thatctxis never used.Two consequences:
- Lines 292-483 build
boundary_positions,seed_*,main_slot_mapping,idx_slot_mapping,final_*_state_mapping,dense_cmp_prefix, anddense_idx_prefixthat no spec consumes. Lines 896-955 then recompute the equivalent leaf slot mappings with the same_lower_rowrule. The two derivations can drift apart, and only the second one reaches the kernel.validate_cp_indexer_capacityat line 289 runs only as a side effect of this otherwise-unused function. If someone removes the dead branch, the capacity guard disappears silently.Reduce
_build_metadata_tensorsto the tables and context thatbuild_tensor_specsactually consumes, and callvalidate_cp_indexer_capacityfrombuild_tensor_specsdirectly.Also applies to: 643-643
🤖 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/prefill_csa_cp.py` around lines 271 - 289, Reduce _build_metadata_tensors to return only the metadata tables and context entries consumed by build_tensor_specs, removing the unused boundary, seed, slot-mapping, state-mapping, and dense-prefix calculations. Move validate_cp_indexer_capacity out of _build_metadata_tensors and invoke it directly from build_tensor_specs, preserving the existing capacity validation behavior independently of metadata construction.Source: Linters/SAST tools
models/deepseek/v4-flash/prefill_exchange_cp.py (1)
478-623: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy lift
_prefill_cp_sparse_stagestill duplicates_cp_swa_stage_sources.This helper and
_cp_swa_stage_sourcesinprefill_swa_cp.pylines 294-379 implement the same persistent/predecessor/current overlay lowering. The gather loop, the overlay index derivation, and the bias build are line-for-line equivalent. This helper adds the compressed-index branch at lines 564-577 and the valid-mask build at lines 603-621; the SWA copy omits both.The PR objective states that exchange logic is consolidated into
prefill_exchange_cp.py, so the SWA copy is an outlier. The SWA path can call_prefill_cp_sparse_stagewithIDX_TOPK-widecmp_indicesfilled with-1, which makes the compressed branch inert and produces the same valid mask.This is deferable, but it prevents the two overlay lowerings from diverging.
🤖 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/prefill_exchange_cp.py` around lines 478 - 623, Refactor the SWA path to reuse _prefill_cp_sparse_stage instead of maintaining the duplicate _cp_swa_stage_sources implementation in prefill_swa_cp.py. Supply IDX_TOPK-wide cmp_indices initialized to -1 and the required compressed-stage buffers so the compressed branch remains inert, while preserving the existing SWA outputs and overlay behavior.models/deepseek/v4-flash/prefill_indexer.py (1)
379-482: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftConsider extracting the shared query-preparation pipeline.
Lines 379-482 of
_prefill_indexer_cp_score_topkrepeat the body ofprefill_indexerlines 148-253 almost verbatim: theqr_projINT8 GEMM and dequantization, the interleaved RoPE table materialization, the swap-index build, the per-token RoPE rotation, the Hadamard multiply, and the per-row INT8 quantization. Only thename_hintstrings differ.Extract one
@pl.jit.inlinehelper that returnsqr_hadamard_i8andqr_hadamard_scale_dq, and give it aname_prefixargument for the task hints. That keeps the two scoring paths from drifting when the quantization math changes.This is deferable if you prefer to land the CP path first, but track it: the same block already exists in three copies across the CP modules.
🤖 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/prefill_indexer.py` around lines 379 - 482, Extract the duplicated query-preparation pipeline from `_prefill_indexer_cp_score_topk` and `prefill_indexer` into one `@pl.jit.inline` helper that returns `qr_hadamard_i8` and `qr_hadamard_scale_dq`. Move the shared projection/dequantization, RoPE preparation and rotation, Hadamard multiplication, and per-row INT8 quantization into the helper, using a `name_prefix` parameter to generate task-specific `name_hint` values, then replace each duplicated block with the helper call.models/deepseek/v4-flash/decode_sparse_attn_csa.py (4)
52-57: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the duplicate
MAX_SEQ_LENassignment.Line 39 already sets
MAX_SEQ_LEN = M.max_position_embeddings. Line 54 repeats the same assignment inside the CSA masking block, where the preceding comment describes compressed-slot masking. The duplicate suggests a different value is intended here.♻️ Proposed cleanup
# CSA compressed-slot masking (folded in from the CSA orchestrator): raw indexer # topk -> per-token bound floor((pos + 1) / COMPRESS_RATIO). -MAX_SEQ_LEN = M.max_position_embeddings INDEXER_SCORE_LEN = MAX_SEQ_LEN // 4🤖 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/decode_sparse_attn_csa.py` around lines 52 - 57, Remove the duplicate MAX_SEQ_LEN assignment from the CSA compressed-slot masking block, while retaining the existing assignment near line 39. Keep INDEXER_SCORE_LEN and the remaining CSA constants using the shared MAX_SEQ_LEN value.
690-710: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename the loop variable to avoid shadowing the masked tensor.
rawis bound at Line 670 to the masked indexer tensor. Both loops here rebindrawto a plainint. The tensor is not read after Line 673, so the result is correct, but the reuse makes the masking block harder to follow.♻️ Proposed rename
- for raw in window_swa_indices[t].tolist(): - slot = int(raw) + for win_slot_raw in window_swa_indices[t].tolist(): + slot = int(win_slot_raw) @@ - for raw in cmp_sparse_indices[t].tolist(): - if raw < 0: + for cmp_slot_raw in cmp_sparse_indices[t].tolist(): + if cmp_slot_raw < 0: kv_rows.append(torch.zeros(HEAD_DIM, dtype=ori_kv.dtype)) valid.append(False) continue - cmp_slot = int(raw) + cmp_slot = int(cmp_slot_raw)🤖 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/decode_sparse_attn_csa.py` around lines 690 - 710, Rename the loop variable raw in both index-iteration loops to distinct descriptive names, updating each corresponding reference while preserving the existing index handling and masking behavior.
542-545: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueHoist the
T_PADzero-fill out of the token loop.This block sits inside
for qt in pl.pipeline(0, T, QUANT_TOKEN_TILE, ...). It writes the same constant zeros to the same padding rows on every token tile, so it repeatsT // QUANT_TOKEN_TILEtimes. The result is correct, but the work is redundant. The padding rows do not depend onqt.♻️ Proposed refactor
o_r_i8_pad = pl.assemble(o_r_i8_pad, oq_i8, [qt, col_g]) - if T_PAD > T: - zero_half = pl.full([T_PAD - T, O_LORA], dtype=pl.FP16, value=0.0) - zero_i8 = pl.cast(zero_half, target_type=pl.INT8, mode="trunc") - o_r_i8_pad = pl.assemble(o_r_i8_pad, zero_i8, [T, col_g]) + if T_PAD > T: + zero_half = pl.full([T_PAD - T, O_LORA], dtype=pl.FP16, value=0.0) + zero_i8 = pl.cast(zero_half, target_type=pl.INT8, mode="trunc") + o_r_i8_pad = pl.assemble(o_r_i8_pad, zero_i8, [T, col_g])🤖 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/decode_sparse_attn_csa.py` around lines 542 - 545, Move the T_PAD padding initialization out of the for qt token-loop and perform it once before the loop, preserving the existing zero_half, zero_i8, and o_r_i8_pad assembly behavior for T_PAD > T. Keep token-dependent processing inside the loop and ensure the padding rows remain unchanged for every tile.
316-337: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider bulk-run detection for the window gather.
This loop issues one
pl.gather_rowper sparse-K row, 128 rows per work item. The sibling kernels avoid that for the window prefix.decode_sparse_attn_swa.sparse_attn_swaanddecode_sparse_attn_hca.sparse_attn_hcaprobe the first and last slot of aGATHER_RUNsub-tile, and copy the whole run as one bulk transfer when the endpoints areGATHER_RUN - 1apart. Window slots are contiguous in the common case, so the same probe applies here.This changes scheduling only, not results. Treat it as deferrable if the current profile is acceptable.
🤖 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/decode_sparse_attn_csa.py` around lines 316 - 337, Optimize the window-prefix path in the loop around qk_kv by detecting contiguous GATHER_RUN sub-tiles, probing the first and last window slots as in decode_sparse_attn_swa.sparse_attn_swa and decode_sparse_attn_hca.sparse_attn_hca, and using one bulk gather when their separation is GATHER_RUN - 1. Preserve the existing per-row gather and fallback behavior for non-contiguous or invalid slots, without changing results for the comparison-sparse path.models/deepseek/v4-flash/decode_swa.py (1)
417-427: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueCache the SWA metadata to avoid recomputing the nested Python loop.
init_swa_indicesandinit_swa_lenseach callinit_swa_metadata, which runsswa_indices_and_lens. That helper loops overT * WINelements with per-element.item()calls. It also rebuildsinit_start_posandinit_block_tableon every call. The values are deterministic, so the results stay consistent, but the fixture pays the cost twice.♻️ Proposed refactor
- def init_swa_metadata(): - return swa_indices_and_lens( - position_ids_from_starts(init_start_pos(), seq=S), - init_block_table(), - block_size=BLOCK_SIZE, - window=WIN, - ) + swa_metadata_cache = {} + + def init_swa_metadata(): + if "value" not in swa_metadata_cache: + swa_metadata_cache["value"] = swa_indices_and_lens( + position_ids_from_starts(init_start_pos(), seq=S), + init_block_table(), + block_size=BLOCK_SIZE, + window=WIN, + ) + return swa_metadata_cache["value"]🤖 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/decode_swa.py` around lines 417 - 427, Cache the result of init_swa_metadata so swa_indices_and_lens, init_start_pos, and init_block_table execute only once. Update init_swa_indices and init_swa_lens to reuse the cached metadata while preserving their respective contiguous outputs.
🤖 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 `@docs/models/deepseek.md`:
- Line 95: Update the entry-point table in the DeepSeek documentation to replace
the old *_attention_csa.py, *_attention_hca.py, and *_attention_swa.py
references with the renamed mode-first modules, including the applicable _cp
variants, while retaining prefill_csa.py where already correct.
In `@models/deepseek/v4-flash/decode_hca.py`:
- Around line 225-241: Update the top-k token iteration around topk_all so the
final partial HCA_TOPK_TOKEN_TILE is included rather than relying on a grid that
truncates the tail. Remove the ineffective topk_t < T guard or replace the loop
bounds with a ceil-divided token-block count, and ensure every token index below
T initializes its HCA_CMP_TOPK entries to valid compressed slots or -1
consistently with sparse_attn_hca expectations.
- Around line 13-14: Update the companion-file references in the module
docstring to use the renamed mode-first files decode_swa.py and decode_csa.py,
replacing the obsolete attention_swa.py and attention_csa_draft.py names.
- Around line 195-202: Add a module-level divisibility assertion for
HCA_WB_TOKEN_TILE, matching the guard style used in decode_sparse_attn_csa.py,
to ensure the configured T value is an exact multiple of HCA_WB_TOKEN_TILE
before the writeback loop runs. Preserve the existing writeback logic in the
hca_cache_writeback loop.
- Around line 273-276: Update attention_hca_test’s compressor-cache handling:
declare compress_state and cmp_kv as pl.InOut, and mark both corresponding
TensorSpecs with is_output=True and output comparison. Ensure
compressor_ratio128 and golden_compressor updates to these buffers are
validated.
- Around line 84-90: Replace HCA_CMP_TOPK usage with HCA_TOPK_LIMIT for the HCA
top-k bound, and add an assertion that HCA_TOPK_LIMIT is at least
get_standalone_cmp_valid(128). Do not reuse HCA_SPARSE_CMP_TOPK for this HCA
limit.
In `@models/deepseek/v4-flash/decode_sparse_attn_csa.py`:
- Line 743: Update the zip call in the loop over block_mi, block_li, and
block_oi to pass strict=True, explicitly validating that all three iterables
have equal lengths while preserving the existing iteration behavior.
- Around line 256-264: Add an assertion alongside the existing tiling-invariant
checks in the module initialization/configuration validation to require WIN ==
ATTN_K_TILE. Keep the block-validity mapping in the loops using its current
offsets, and ensure invalid configurations fail before qk_pv executes.
In `@models/deepseek/v4-flash/decode_swa.py`:
- Around line 14-15: Update the module docstring’s companion file references to
use decode_csa.py and decode_hca.py instead of the renamed
attention_csa_draft.py and attention_hca_draft.py entries, preserving the
existing ratio annotations.
In `@models/deepseek/v4-flash/prefill_csa_cp.py`:
- Line 1973: Update the sparse_bias initialization in the prefill sparse
attention path to import FP32_NEG_INF from config and use it instead of the
literal -3.0e38, keeping the existing tensor shape, dtype, and initialization
behavior unchanged.
In `@models/deepseek/v4-flash/prefill_csa.py`:
- Around line 231-237: The sparse-attention launch following
`_prefill_compressor_ratio4_with_completion` must depend on the compressor write
completion before reading `cmp_kv`. Update the
`_prefill_sparse_attn_with_block_mask` call or its launch dependency arguments
to include `compressor_completion[0]`, preserving the same CP ordering behavior
used in `prefill_csa_cp.py`.
In `@models/deepseek/v4-flash/prefill_hca_cp.py`:
- Around line 784-807: Update the HCA compression flow around
prefill_compressor_ratio128 and cp_hca_pack_compact to collect the compressor
completion IDs for each launched leaf/state write, then pass or commission those
completions through the compact payload dependency path before
cp_hca_pack_compact reads leaf_cmp_flat and scratch_state_flat. Preserve the
existing buffer slicing and packing behavior while ensuring compact packing
waits on all compressor roots, matching the CSA ratio-4 caller’s dependency
handling.
In `@models/deepseek/v4-flash/prefill_swa_cp.py`:
- Around line 306-310: Populate valid_block_mask in prefill_cp_swa using the
same mask-construction logic as _prefill_cp_sparse_stage before invoking
_prefill_sparse_attn_math. Ensure every staged attention block row is enabled
according to the SWA CP staging inputs, rather than leaving nonzero rows masked
by the zero-initialized valid_mask.
---
Outside diff comments:
In `@models/deepseek/v4-flash/prefill_indexer_compressor.py`:
- Around line 347-393: Add an explicit dependency on cache_write_tid to the
prefill_idx_c4_scale_scatter task declaration so it waits for the producer that
writes scale_scratch. Preserve the existing scale scatter logic and dependency
behavior for all other tasks.
---
Nitpick comments:
In `@models/deepseek/v4-flash/decode_sparse_attn_csa.py`:
- Around line 52-57: Remove the duplicate MAX_SEQ_LEN assignment from the CSA
compressed-slot masking block, while retaining the existing assignment near line
39. Keep INDEXER_SCORE_LEN and the remaining CSA constants using the shared
MAX_SEQ_LEN value.
- Around line 690-710: Rename the loop variable raw in both index-iteration
loops to distinct descriptive names, updating each corresponding reference while
preserving the existing index handling and masking behavior.
- Around line 542-545: Move the T_PAD padding initialization out of the for qt
token-loop and perform it once before the loop, preserving the existing
zero_half, zero_i8, and o_r_i8_pad assembly behavior for T_PAD > T. Keep
token-dependent processing inside the loop and ensure the padding rows remain
unchanged for every tile.
- Around line 316-337: Optimize the window-prefix path in the loop around qk_kv
by detecting contiguous GATHER_RUN sub-tiles, probing the first and last window
slots as in decode_sparse_attn_swa.sparse_attn_swa and
decode_sparse_attn_hca.sparse_attn_hca, and using one bulk gather when their
separation is GATHER_RUN - 1. Preserve the existing per-row gather and fallback
behavior for non-contiguous or invalid slots, without changing results for the
comparison-sparse path.
In `@models/deepseek/v4-flash/decode_swa.py`:
- Around line 417-427: Cache the result of init_swa_metadata so
swa_indices_and_lens, init_start_pos, and init_block_table execute only once.
Update init_swa_indices and init_swa_lens to reuse the cached metadata while
preserving their respective contiguous outputs.
In `@models/deepseek/v4-flash/prefill_csa_cp.py`:
- Around line 271-289: Reduce _build_metadata_tensors to return only the
metadata tables and context entries consumed by build_tensor_specs, removing the
unused boundary, seed, slot-mapping, state-mapping, and dense-prefix
calculations. Move validate_cp_indexer_capacity out of _build_metadata_tensors
and invoke it directly from build_tensor_specs, preserving the existing capacity
validation behavior independently of metadata construction.
In `@models/deepseek/v4-flash/prefill_csa.py`:
- Around line 251-259: Update the assignment receiving outputs from
prefill_indexer so the intentionally unused idx_kv_cache_out and
idx_kv_scale_out bindings use underscore-prefixed names, while preserving the
remaining outputs and call behavior.
In `@models/deepseek/v4-flash/prefill_exchange_cp.py`:
- Around line 478-623: Refactor the SWA path to reuse _prefill_cp_sparse_stage
instead of maintaining the duplicate _cp_swa_stage_sources implementation in
prefill_swa_cp.py. Supply IDX_TOPK-wide cmp_indices initialized to -1 and the
required compressed-stage buffers so the compressed branch remains inert, while
preserving the existing SWA outputs and overlay behavior.
In `@models/deepseek/v4-flash/prefill_hca_cp.py`:
- Around line 647-662: In cp_hca_tail_assemble, remove the redundant if/else
around source and assign it directly using part * MAX_SEGMENT_TILES * TAIL_ROWS
+ tail_offset. Apply the same simplification to the corresponding tail assembly
logic in prefill_swa_cp.py, preserving the existing contiguous tile layout and
indexing behavior.
- Around line 1313-1314: Derive the shared segment geometry once before the
metadata construction in the fixture setup, then pass that geometry into both
build_hca_metadata and _build_raw_attention_metadata instead of recomputing
prefix, span, lengths, starts, owners, and query_positions independently. Update
both builder signatures and their callers, and ensure device_values uses fields
from this single consistent geometry source.
In `@models/deepseek/v4-flash/prefill_indexer.py`:
- Around line 379-482: Extract the duplicated query-preparation pipeline from
`_prefill_indexer_cp_score_topk` and `prefill_indexer` into one `@pl.jit.inline`
helper that returns `qr_hadamard_i8` and `qr_hadamard_scale_dq`. Move the shared
projection/dequantization, RoPE preparation and rotation, Hadamard
multiplication, and per-row INT8 quantization into the helper, using a
`name_prefix` parameter to generate task-specific `name_hint` values, then
replace each duplicated block with the helper call.
In `@models/deepseek/v4-flash/prefill_swa_cp.py`:
- Around line 502-510: Rename the local scalar `tail_start` in `prefill_cp_swa`
to `tail_start0`, and update its use when computing `tail_offset`; preserve the
existing tail-copy behavior while avoiding shadowing the module-level
`tail_start` helper.
- Around line 642-644: Update build_tensor_specs to reject any cp_size that
differs from the import-time CP_SIZE before building metadata, set
golden_prefill_cp_swa._ctx within the builder, and return only specs. Adjust the
__main__ block to use the builder’s new return contract without externally
installing _ctx.
In `@models/deepseek/v4-flash/prefill_zigzag_cp.py`:
- Around line 33-43: Export and reuse the existing _parse_static_int helper from
prefill_zigzag_cp.py in prefill_swa_cp.py, removing the duplicate implementation
and deriving CP_SIZE through the shared parser (or importing CP_SIZE directly).
Ensure both modules use one consistent --cp argument parsing rule.
🪄 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 Plus
Run ID: 3f6cfc98-1646-4186-8e55-061d272ea038
📒 Files selected for processing (27)
docs/debug-and-tune/debugging.mddocs/debug-and-tune/performance-tuning.mddocs/models/deepseek.mddocs/run-and-validate/compile-runtime-workflow.mdmodels/deepseek/v4-flash/decode_csa.pymodels/deepseek/v4-flash/decode_fwd.pymodels/deepseek/v4-flash/decode_hca.pymodels/deepseek/v4-flash/decode_indexer.pymodels/deepseek/v4-flash/decode_layer.pymodels/deepseek/v4-flash/decode_mtp.pymodels/deepseek/v4-flash/decode_sparse_attn_csa.pymodels/deepseek/v4-flash/decode_swa.pymodels/deepseek/v4-flash/prefill_compressor_ratio4.pymodels/deepseek/v4-flash/prefill_csa.pymodels/deepseek/v4-flash/prefill_csa_cp.pymodels/deepseek/v4-flash/prefill_exchange_cp.pymodels/deepseek/v4-flash/prefill_fwd.pymodels/deepseek/v4-flash/prefill_hca.pymodels/deepseek/v4-flash/prefill_hca_cp.pymodels/deepseek/v4-flash/prefill_indexer.pymodels/deepseek/v4-flash/prefill_indexer_compressor.pymodels/deepseek/v4-flash/prefill_layer.pymodels/deepseek/v4-flash/prefill_mtp.pymodels/deepseek/v4-flash/prefill_sparse_attn.pymodels/deepseek/v4-flash/prefill_swa.pymodels/deepseek/v4-flash/prefill_swa_cp.pymodels/deepseek/v4-flash/prefill_zigzag_cp.py
🛑 Comments failed to post (9)
models/deepseek/v4-flash/decode_hca.py (5)
13-14: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the companion file names after the rename.
The docstring points to
attention_swa.pyandattention_csa_draft.py. This PR renames the attention entry files to mode-first names, so these paths no longer exist. Usedecode_swa.pyanddecode_csa.py.📝 Proposed doc fix
-Companion files: attention_swa.py (ratio=0) - attention_csa_draft.py (ratio=4).""" +Companion files: decode_swa.py (ratio=0) + decode_csa.py (ratio=4)."""As per coding guidelines "Keep public documentation and examples aligned when behavior changes".
📝 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.Companion files: decode_swa.py (ratio=0) decode_csa.py (ratio=4)."""🤖 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/decode_hca.py` around lines 13 - 14, Update the companion-file references in the module docstring to use the renamed mode-first files decode_swa.py and decode_csa.py, replacing the obsolete attention_swa.py and attention_csa_draft.py names.Source: Coding guidelines
84-90: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Description: Compare the HCA compressed topk bound with the sparse module's CMP_TOPK. set -euo pipefail fd -t f 'decode_sparse_attn_hca.py' | while IFS= read -r file; do echo "== $file" rg -n -C 3 '^CMP_TOPK|^IDX_TOPK|^TOPK\b|^WIN\b|^PADDED_TOPK' "$file" done fd -t f 'config.py' --full-path 'v4-flash' | while IFS= read -r file; do rg -n 'index_topk|max_position_embeddings|sliding_window' "$file" doneRepository: hw-native-sys/pypto-lib
Length of output: 2951
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== locate v4-flash files ==" fd -t f 'decode_hca.py|decode_sparse_attn_hca.py|config.py|.*\.json$' models/deepseek/v4-flash -x sh -c 'echo "-- $1"; wc -l "$1"' sh {} echo "== relevant decode_hca.py sections ==" sed -n '1,120p' models/deepseek/v4-flash/decode_hca.py sed -n '220,255p' models/deepseek/v4-flash/decode_hca.py echo "== search COMPRESS_RATIO definitions/usages ==" rg -n -C 2 'COMPRESS_RATIO|DEFAULT_COMPRESS_RATIO|index_topk|max_position_embeddings|sliding_window|MAX_SEQ_LEN|CMP_TOPK|HCA_TOPK_LIMIT|HCA_CMP_TOPK|topk_all' models/deepseek/v4-flash echo "== compute constants from source text (no repo code run) ==" python3 - <<'PY' import pathlib, re fp = pathlib.Path("models/deepseek/v4-flash/decode_sparse_attn_hca.py") text = fp.read_text() m = dict(re.findall(r'^(?:PADDED_)?TOPK\s*=\s*(.+)$|^(?:DEFAULT_)?COMPRESS_RATIO\s*=\s*(.+)$|^M\.(.+)$', text, re.M)) # extract expressions and evaluate only if safe numeric-like def safe(expr): try: return eval(expr, {"__builtins__":{}}, { "MAX_SEQ_LEN": 16384, "sliding_window": 512, "ATTN_K_TILE": next(int(s) for s in re.findall(r'\bATTN_K_TILE\b', text)), "BLOCK_SIZE": 16, "GATHER_SEGS": 2, }) except Exception as e: return f"<eval_failed: {e!r}>" attrs = ["MAX_SEQ_LEN", "sliding_window", "DEFAULT_COMPRESS_RATIO", "ATTN_K_TILE", "BLOCK_SIZE", "GATHER_SEGS"] for a in attrs: print(f"M.{a} = {safe(a)}") for name in ["PADDED_TOPK", "TOPK", "CMP_TOPK", "SPARSE_BLOCKS"]: print(f"{name} = {m.get(name)} -> {safe(m.get(name))}") PYRepository: hw-native-sys/pypto-lib
Length of output: 50380
🏁 Script executed:
#!/bin/bash #!/bin/bash set -euo pipefail echo "== targeted constants at v4-flash decode files ==" for f in models/deepseek/v4-flash/decode_hca.py models/deepseek/v4-flash/decode_sparse_attn_hca.py; do echo "-- $f" rg -n -C 3 '^(COMPRESS_RATIO|DEFAULT_COMPRESS_ratio|HCA_TOPK_LIMIT|HCA_CMP_TOPK|AT...|ATTN_K|BLOCK_SIZE|PADDED_TOPK|TOPK)$|^(CMP_)?(TOPK|CAPACITY|BLOCKS|TOPK_LIMIT).*=|from .* import' "$f" done echo "== exact targeted definitions only ==" python3 - <<'PY' import pathlib, re files = [ pathlib.Path("models/deepseek/v4-flash/decode_hca.py"), pathlib.Path("models/deepseek/v4-flash/decode_sparse_attn_hca.py"), ] for p in files: print(f"-- {p}") text = p.read_text() for pat in [ r'(?m)^COMPRESS_RATIO\s*=\s*(.+)$', r'(?m)^DEFAULT_COMPRESS_RATIO\s*=\s*(.+)$', r'(?m)^HCA_TOPK_LIMIT\s*=\s*(.+)$', r'(?m)^HCA_CMP_TOPK\s*=\s*(.+)$', r'(?m)^CMP_TOPK\s*=\s*(.+)$', r'(?m)^MAYBE_BLOCK_SIZE\s*=\s*(.+)$', r'(?m)^AT..._K_TILE\s*=\s*(.+)$', r'(?m)^TOPK\s*=\s*(.+)$', r'(?m)^PADDED_TOPK\s*=\s*(.+)$', ]: m = re.search(pat, text) if m: print(f"{pat.split('s',1)[0].strip()[1:-1]} = {m.group(1)}") PYRepository: hw-native-sys/pypto-lib
Length of output: 4585
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== decode_hca.py relevant constants ==" sed -n '60,95p' models/deepseek/v4-flash/decode_hca.py echo "== decode_sparse_attn_hca.py constants and sparse-attn loop ==" rg -n -C 4 '^(AT..._K_TILE|BLOCK_SIZE|MAYBE_|ATTN_|DEFAULT_|CMP_CAPACITY|CMP_TOPK|TOPK|PADDED_TOPK)$|^def sparse_attn_hca|^CMP_TOPK' models/deepseek/v4-flash/decode_sparse_attn_hca.py echo "== all topk_usages and loop boundaries in sparse model ==" rg -n -C 3 'topk_all|range\(|valid_comp|CMP_TOPK|TOPK|window_start|position|score_pos|kv_pos' models/deepseek/v4-flash/decode_sparse_attn_hca.py echo "== numeric probe from source constants ==" python3 - <<'PY' import pathlib, re from collections import OrderedDict files = [ ("models/deepseek/v4-flash/decode_hca.py", ["COMPRESS_RATIO", "DEFAULT_COMPRESS_RATIO", "HCA_TOPK_LIMIT", "HCA_CMP_TOPK", "PADDED_TOPK"]), ("models/deepseek/v4-flash/decode_sparse_attn_hca.py", ["WIN", "PADDED_TOPK", "ATTN_K_TILE", "BLOCK_SIZE", "MAYBE_BLOCK_SIZE", "COMPACTION...", "DEFAULT_COMPRESS_RATIO", "CMP_CAPACITY", "CMP_TOPK", "MAX_SEQ_LEN", "TOPK", "TOPK_FULL", "SPARSE_BLOCKS", "GATHER", "GATHER_SEGS"]), ] for path, wanted in files: print(f"-- {path} wanted") text = pathlib.Path(path).read_text() for pat in wanted: m = re.search(fr'(?m)^{re.escape(pat)}\s*=\s*(.+)$', text) if m: print(f"{pat}={m.group(1)}") PYRepository: hw-native-sys/pypto-lib
Length of output: 19052
Replace
HCA_CMP_TOPKwithHCA_TOPK_LIMIT.
HCA_TOPK_LIMITisMAX_SEQ_LEN // 128(128), butHCA_CMP_TOPKfollowsget_standalone_cmp_valid()and can be smaller for non-128 ratios. UseHCA_TOPK_LIMITfor the HCA top-k bound and assertHCA_TOPK_LIMIT >= get_standalone_cmp_valid(128)instead of reusing the generic sparseCMP_TOPK.🤖 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/decode_hca.py` around lines 84 - 90, Replace HCA_CMP_TOPK usage with HCA_TOPK_LIMIT for the HCA top-k bound, and add an assertion that HCA_TOPK_LIMIT is at least get_standalone_cmp_valid(128). Do not reuse HCA_SPARSE_CMP_TOPK for this HCA limit.
195-202: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Add a divisibility assertion for
HCA_WB_TOKEN_TILE.The grid is
T // HCA_WB_TOKEN_TILE. IfTis not a multiple ofHCA_WB_TOKEN_TILE, the lastT % HCA_WB_TOKEN_TILEtokens are never written intokv_cache. The kernel produces no error; the cache silently keeps stale rows and attention reads wrong KV.T = DECODE_BATCH * DECODE_SEQcomes from config, so a config change can break this without warning.decode_sparse_attn_csa.pyguards the same class of assumption with module-level asserts (Lines 141-154).🛡️ Proposed guard
HCA_TOPK_TOKEN_TILE = 8 # tokens per cache-window topk SPMD block HCA_WB_TOKEN_TILE = 8 # tokens per cache-writeback SPMD block + +assert T % HCA_WB_TOKEN_TILE == 0, "cache writeback grid must cover every token" +assert T % HCA_TOPK_TOKEN_TILE == 0, "cache topk grid must cover every token"🤖 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/decode_hca.py` around lines 195 - 202, Add a module-level divisibility assertion for HCA_WB_TOKEN_TILE, matching the guard style used in decode_sparse_attn_csa.py, to ensure the configured T value is an exact multiple of HCA_WB_TOKEN_TILE before the writeback loop runs. Preserve the existing writeback logic in the hca_cache_writeback loop.
225-241: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
The
topk_t < Tguard is unreachable, and the tail tokens are unreachable too.
topk_t0is at most(T // HCA_TOPK_TOKEN_TILE - 1) * HCA_TOPK_TOKEN_TILE, andtopk_dt < HCA_TOPK_TOKEN_TILE, sotopk_tis always belowT. The condition at Line 229 never fails. The real gap is the opposite: whenTis not a multiple ofHCA_TOPK_TOKEN_TILE, the tail tokens keep the uninitializedtopk_allcontents, whichsparse_attn_hcathen reads as compressed slots. This shares the root cause with the writeback grid at Lines 195-202.🤖 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/decode_hca.py` around lines 225 - 241, Update the top-k token iteration around topk_all so the final partial HCA_TOPK_TOKEN_TILE is included rather than relying on a grid that truncates the tail. Remove the ineffective topk_t < T guard or replace the loop bounds with a ceil-divided token-block count, and ensure every token index below T initializes its HCA_CMP_TOPK entries to valid compressed slots or -1 consistently with sparse_attn_hca expectations.
273-276: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash # Description: Compare InOut/is_output marking for compressor-written caches across v4-flash modules. set -euo pipefail rg -n -C 2 'cmp_kv|compress_state' --glob 'models/deepseek/v4-flash/*.py' | rg -n 'InOut|is_output' -C 2Repository: hw-native-sys/pypto-lib
Length of output: 161
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== candidates ==" git ls-files 'models/deepseek/v4-flash/*.py' | sed -n '1,200p' echo "== target file line count ==" wc -l models/deepseek/v4-flash/decode_hca.py echo "== decode_hca relevant sections ==" sed -n '240,310p' models/deepseek/v4-flash/decode_hca.py sed -n '630,660p' models/deepseek/v4-flash/decode_hca.pyRepository: hw-native-sys/pypto-lib
Length of output: 8292
🏁 Script executed:
#!/bin/bash set -euo pipefail echo "== decode_compressor_ratio128 outline ==" ast-grep outline models/deepseek/v4-flash/decode_compressor_ratio128.py --view compact || true echo "== decode_compressor_ratio128 relevant writes ==" rg -n -C 3 'compress_state|cmp_kv|compressor_ratio128|def golden|is_input|InOut|In|Out' models/deepseek/v4-flash/decode_compressor_ratio128.py echo "== prefill compressor tensor specs ==" rg -n -C 2 'compress_state|cmp_kv|TensorSpec' models/deepseek/v4-flash/prefill_compressor_ratio128.py echo "== compare svmlib/prefill/sparse HCA signatures and specs ==" for f in \ models/deepseek/v4-flash/prefill_hca.py \ models/deepseek/v4-flash/decode_swa.py \ models/deepseek/v4-flash/decode_sparse_attn_hca.py \ models/deepseek/v4-flash/decode_metadata.py do if [ -f "$f" ]; then echo "--- $f ---" rg -n -C 2 'cmp_kv|compress_state|pl\.InOut|TensorSpec\("compress_state"|TensorSpec\("cmp_kv"|is_output=True' "$f" || true fi doneRepository: hw-native-sys/pypto-lib
Length of output: 31902
Mark
compress_stateandcmp_kvas outputs inattention_hca_test.
compressor_ratio128writescompress_staterows and the pooledcmp_kv_cache;golden_compressordoes the same. In this harness, pass them aspl.InOut, and addis_output=Trueplus comparison for bothcompress_stateandcmp_kvin the TensorSpecs so the compressor cache updates are exercised.🤖 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/decode_hca.py` around lines 273 - 276, Update attention_hca_test’s compressor-cache handling: declare compress_state and cmp_kv as pl.InOut, and mark both corresponding TensorSpecs with is_output=True and output comparison. Ensure compressor_ratio128 and golden_compressor updates to these buffers are validated.models/deepseek/v4-flash/decode_sparse_attn_csa.py (2)
256-264: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
The block-validity mapping assumes
WIN == ATTN_K_TILE; add an assertion.Sparse block
sbcovers sparse-K columns[sb*ATTN_K_TILE, (sb+1)*ATTN_K_TILE). The code maps blocksb >= 1to compressed-mask columns[(sb-1)*ATTN_K_TILE, sb*ATTN_K_TILE), and it treats block 0 as pure window. Both mappings hold only whenWINequalsATTN_K_TILEexactly.If
WIN > ATTN_K_TILE, block 1 still contains window columns but its validity comes from the compressed mask. IfWIN < ATTN_K_TILE, block 0 already contains compressed columns and block 1 reads the wrong compressed offset. In both casesqk_pvcan take the empty branch at Line 370 for a block that holds valid KV, writemi = -3.0e38andli = 0, and drop that attention mass with no error.The module already asserts the other tiling invariants at Lines 141-154. Add this one.
🛡️ Proposed guard
assert WIN <= TOPK <= TOPK_FULL, f"TOPK ({TOPK}) must be in [WIN={WIN}, TOPK_FULL={TOPK_FULL}]" +# Block 0 is treated as pure window and block sb>=1 maps to compressed-mask +# columns [(sb-1)*ATTN_K_TILE, sb*ATTN_K_TILE); both hold only at this equality. +assert WIN == ATTN_K_TILE, f"valid_block_mask assumes WIN ({WIN}) == ATTN_K_TILE ({ATTN_K_TILE})"🤖 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/decode_sparse_attn_csa.py` around lines 256 - 264, Add an assertion alongside the existing tiling-invariant checks in the module initialization/configuration validation to require WIN == ATTN_K_TILE. Keep the block-validity mapping in the loops using its current offsets, and ensure invalid configurations fail before qk_pv executes.
743-743: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Add
strict=Truetozip().Ruff reports B905 on this call. The three lists always have the same length, so the behavior does not change today, but
strict=Truemakes that invariant explicit and catches a future mismatch.🔧 Proposed fix
- for mi_cur, li_cur, oi_cur in zip(block_mi[1:], block_li[1:], block_oi[1:]): + for mi_cur, li_cur, oi_cur in zip(block_mi[1:], block_li[1:], block_oi[1:], strict=True):📝 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.for mi_cur, li_cur, oi_cur in zip(block_mi[1:], block_li[1:], block_oi[1:], strict=True):🧰 Tools
🪛 Ruff (0.16.0)
[warning] 743-743:
zip()without an explicitstrict=parameterAdd explicit value for parameter
strict=(B905)
🤖 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/decode_sparse_attn_csa.py` at line 743, Update the zip call in the loop over block_mi, block_li, and block_oi to pass strict=True, explicitly validating that all three iterables have equal lengths while preserving the existing iteration behavior.Source: Linters/SAST tools
models/deepseek/v4-flash/decode_swa.py (1)
14-15: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the companion file names after the rename.
This PR renames the attention entry files to mode-first names. The docstring still points to
attention_csa_draft.pyandattention_hca_draft.py, which no longer exist. Usedecode_csa.pyanddecode_hca.py.📝 Proposed doc fix
-Companion files: attention_csa_draft.py (ratio=4) - attention_hca_draft.py (ratio=128).""" +Companion files: decode_csa.py (ratio=4) + decode_hca.py (ratio=128)."""As per coding guidelines "Keep public documentation and examples aligned when behavior changes".
📝 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.Companion files: decode_csa.py (ratio=4) decode_hca.py (ratio=128)."""🤖 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/decode_swa.py` around lines 14 - 15, Update the module docstring’s companion file references to use decode_csa.py and decode_hca.py instead of the renamed attention_csa_draft.py and attention_hca_draft.py entries, preserving the existing ratio annotations.Source: Coding guidelines
models/deepseek/v4-flash/prefill_csa_cp.py (1)
1973-1973: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Use the shared
FP32_NEG_INFsentinel forsparse_bias.This line initializes
sparse_biasto the literal-3.0e38.prefill_hca_cp.pyline 924 andprefill_swa_cp.pyline 452 initialize the same buffer withFP32_NEG_INFfromconfig._prefill_cp_sparse_stagealso writesFP32_NEG_INFinto the padding columns atprefill_exchange_cp.pylines 597-601.If the two values differ, masked and padded columns in this path carry different sentinels than the other CP paths, which changes the
pl.row_maxandpl.expbehavior in_prefill_sparse_attn_math. ImportFP32_NEG_INFfromconfigand use it here.♻️ Proposed fix
from config import ( BLOCK_SIZE, CSA_INNER_STATE_PHYSICAL_BLOCKS, CSA_STATE_PHYSICAL_BLOCKS, FLASH as M, + FP32_NEG_INF, IDX_CACHE_MAX_BLOCKS,- sparse_bias = pl.create_tensor([LOCAL_ROWS, PREFILL_SPARSE_PAD], dtype=pl.FP32, init_value=-3.0e38) + sparse_bias = pl.create_tensor([LOCAL_ROWS, PREFILL_SPARSE_PAD], dtype=pl.FP32, init_value=FP32_NEG_INF)📝 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.sparse_bias = pl.create_tensor([LOCAL_ROWS, PREFILL_SPARSE_PAD], dtype=pl.FP32, init_value=FP32_NEG_INF)🤖 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/prefill_csa_cp.py` at line 1973, Update the sparse_bias initialization in the prefill sparse attention path to import FP32_NEG_INF from config and use it instead of the literal -3.0e38, keeping the existing tensor shape, dtype, and initialization behavior unchanged.
1 participant
Summary
Add correctness-first context-parallel prefill support for DeepSeek V4 Flash.
prefill_exchange_cp.py.Design
Shared baseline seams
prefill_compressor_ratio4.py: completion-aware CP composition without changing the standalone path.prefill_indexer_compressor.py: completion-aware inner-compressor composition.prefill_indexer.py: CP score/top-k composition for receiver-local index history.prefill_sparse_attn.py: fixed-row sparse math and source-staging constants used by CP.prefill_csa.py: shared completion/indexer seams while preserving baseline behavior.Latest stabilization
pl.sliceboundaries.InOut, and capture compressor SSA returns before assembling them into the parent tensors.TileTypecommunication sources.Validation
Final CP2/CP4 runtime and golden matrix
task_20260803_051602_231811720024_jit_prefill_cp_zigzag_kv_tail_exchange_test_20260803_051605task_20260803_051636_234358831874_jit_prefill_cp_zigzag_kv_tail_exchange_test_20260803_051639task_20260803_051716_23785312097_jit_prefill_cp_swa_test_20260803_051720task_20260803_051822_258260720944_jit_prefill_cp_swa_test_20260803_051827task_20260803_052256_307288431301_jit_prefill_cp_hca_test_20260803_052301task_20260803_052414_33836153137_jit_prefill_cp_hca_test_20260803_052421task_20260803_053158_121001810760_jit_prefill_cp_csa_test_20260803_053204task_20260803_053344_158081518768_jit_prefill_cp_csa_test_20260803_053350All expected output comparisons passed, including
x_out,kv_cache, and path-specific persistent cache/state outputs. No 507018 or HCCL stall was observed.CP8 bring-up coverage
Shared standalone paths
x_outandkv_cacheLocal checks
git diff --check: PASS.