Enable DeepSeek V4 decode 8k metadata support#688
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds a shared ChangesDeepSeek V4 decode metadata and kernel guard refactor
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant TestBuilder as build_tensor_specs
participant Metadata as decode_metadata
participant Kernel as decode_*_kernel
participant Golden as golden_reference
TestBuilder->>Metadata: resolve_start_positions(start_pos)
Metadata-->>TestBuilder: starts tensor
TestBuilder->>Metadata: position_ids_from_starts / kv_seq_lens_from_starts
TestBuilder->>Metadata: block_table(...)
Metadata-->>TestBuilder: block table
TestBuilder->>Metadata: ori_slot_mapping / compressed_slot_mapping / state_slot_mapping
Metadata-->>TestBuilder: slot mapping tensors
TestBuilder->>Kernel: run kernel with tensors
Kernel->>Kernel: read slot_mapping as i64, check >= 0, cast to INDEX
Kernel->>Kernel: conditional cache/state write
TestBuilder->>Golden: golden_reference with same tensors
Golden->>Golden: guard invalid block ids, zero/-inf defaults
Golden-->>TestBuilder: expected output
TestBuilder->>TestBuilder: compare kernel output vs golden
Possibly related issues
Possibly related PRs
Suggested labels: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a new metadata helper module (decode_metadata.py) to centralize and lower decode fixture metadata, refactoring several decode and attention files to use these unified helpers. The code review feedback highlights significant performance improvement opportunities in the newly added decode_metadata.py module. Specifically, the reviewer points out that several functions rely on nested Python loops and .item() calls on PyTorch tensors, which cause inefficient host-device synchronization. The reviewer recommends fully vectorizing these operations—including position_ids_from_starts, block_table, ori_slot_mapping, compressed_slot_mapping, mask_uncommitted_compressed_boundaries, and state_slot_mapping—using PyTorch broadcasting, torch.gather, and tensor masking to avoid overhead and enhance execution speed.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def position_ids_from_starts(starts: torch.Tensor, *, seq: int = DECODE_SEQ) -> torch.Tensor: | ||
| starts_i64 = starts.to(torch.int64) | ||
| positions = torch.empty((int(starts_i64.numel()), seq), dtype=torch.int32) | ||
| for b in range(int(starts_i64.numel())): | ||
| for s in range(seq): | ||
| positions[b, s] = int(starts_i64[b].item()) + s | ||
| return positions |
There was a problem hiding this comment.
Using nested Python loops and calling .item() on PyTorch tensors inside a loop is highly inefficient and causes host-device synchronization if the tensor is on GPU/NPU. This can be fully vectorized using PyTorch broadcasting, which is much faster and cleaner.
| def position_ids_from_starts(starts: torch.Tensor, *, seq: int = DECODE_SEQ) -> torch.Tensor: | |
| starts_i64 = starts.to(torch.int64) | |
| positions = torch.empty((int(starts_i64.numel()), seq), dtype=torch.int32) | |
| for b in range(int(starts_i64.numel())): | |
| for s in range(seq): | |
| positions[b, s] = int(starts_i64[b].item()) + s | |
| return positions | |
| def position_ids_from_starts(starts: torch.Tensor, *, seq: int = DECODE_SEQ) -> torch.Tensor: | |
| return starts.unsqueeze(1) + torch.arange(seq, dtype=torch.int32, device=starts.device) |
| def block_table( | ||
| *, | ||
| batch: int, | ||
| table_blocks: int, | ||
| physical_blocks: int | None = None, | ||
| permuted: bool = False, | ||
| ) -> torch.Tensor: | ||
| physical_blocks = table_blocks if physical_blocks is None else physical_blocks | ||
| tbl = torch.full((batch, table_blocks), -1, dtype=torch.int32) | ||
| for b in range(batch): | ||
| for j in range(table_blocks): | ||
| phys_j = j % physical_blocks | ||
| if permuted and physical_blocks > 1: | ||
| phys_j = (phys_j * 7 + 3) % physical_blocks | ||
| tbl[b, j] = b * physical_blocks + phys_j | ||
| return tbl |
There was a problem hiding this comment.
This nested loop can be fully vectorized using torch.arange and broadcasting, avoiding Python loop overhead and improving performance.
def block_table(
*,
batch: int,
table_blocks: int,
physical_blocks: int | None = None,
permuted: bool = False,
) -> torch.Tensor:
physical_blocks = table_blocks if physical_blocks is None else physical_blocks
j = torch.arange(table_blocks, dtype=torch.int32)
phys_j = j % physical_blocks
if permuted and physical_blocks > 1:
phys_j = (phys_j * 7 + 3) % physical_blocks
b = torch.arange(batch, dtype=torch.int32).unsqueeze(1)
return b * physical_blocks + phys_j.unsqueeze(0)| def ori_slot_mapping( | ||
| positions: torch.Tensor, | ||
| ori_block_table: torch.Tensor, | ||
| *, | ||
| block_size: int = BLOCK_SIZE, | ||
| window: int = M.sliding_window, | ||
| ) -> torch.Tensor: | ||
| mapping = torch.full(positions.shape, -1, dtype=torch.int64) | ||
| table_i64 = ori_block_table.to(torch.int64) | ||
| positions_i64 = positions.to(torch.int64) | ||
| for b in range(positions_i64.shape[0]): | ||
| for s in range(positions_i64.shape[1]): | ||
| slot = int(positions_i64[b, s].item()) % window | ||
| logical_blk = slot // block_size | ||
| intra = slot % block_size | ||
| blk = int(table_i64[b, logical_blk].item()) | ||
| if blk >= 0: | ||
| mapping[b, s] = blk * block_size + intra | ||
| return mapping |
There was a problem hiding this comment.
Using nested Python loops and .item() calls inside a loop causes host-device synchronization and is extremely slow. This can be fully vectorized using torch.gather and torch.where.
def ori_slot_mapping(
positions: torch.Tensor,
ori_block_table: torch.Tensor,
*,
block_size: int = BLOCK_SIZE,
window: int = M.sliding_window,
) -> torch.Tensor:
positions_i64 = positions.to(torch.int64)
table_i64 = ori_block_table.to(torch.int64)
slot = positions_i64 % window
logical_blk = slot // block_size
intra = slot % block_size
blk = torch.gather(table_i64, 1, logical_blk)
return torch.where(blk >= 0, blk * block_size + intra, -1)| def compressed_slot_mapping( | ||
| positions: torch.Tensor, | ||
| cmp_block_table: torch.Tensor, | ||
| *, | ||
| compress_ratio: int, | ||
| block_size: int = BLOCK_SIZE, | ||
| ) -> torch.Tensor: | ||
| mapping = torch.full(positions.shape, -1, dtype=torch.int64) | ||
| table_i64 = cmp_block_table.to(torch.int64) | ||
| positions_i64 = positions.to(torch.int64) | ||
| for b in range(positions_i64.shape[0]): | ||
| for s in range(positions_i64.shape[1]): | ||
| pos = int(positions_i64[b, s].item()) | ||
| if (pos + 1) % compress_ratio != 0: | ||
| continue | ||
| cache_col = pos // compress_ratio | ||
| logical_blk = cache_col // block_size | ||
| intra = cache_col % block_size | ||
| if logical_blk >= table_i64.shape[1]: | ||
| continue | ||
| blk = int(table_i64[b, logical_blk].item()) | ||
| if blk >= 0: | ||
| mapping[b, s] = blk * block_size + intra | ||
| return mapping |
There was a problem hiding this comment.
This function can be fully vectorized using PyTorch operations like torch.gather, torch.clamp, and torch.where, completely eliminating the nested loops and .item() calls.
def compressed_slot_mapping(
positions: torch.Tensor,
cmp_block_table: torch.Tensor,
*,
compress_ratio: int,
block_size: int = BLOCK_SIZE,
) -> torch.Tensor:
positions_i64 = positions.to(torch.int64)
table_i64 = cmp_block_table.to(torch.int64)
mask = (positions_i64 + 1) % compress_ratio == 0
cache_col = positions_i64 // compress_ratio
logical_blk = cache_col // block_size
intra = cache_col % block_size
in_bounds = logical_blk < table_i64.shape[1]
clamped_blk = torch.clamp(logical_blk, max=table_i64.shape[1] - 1)
blk = torch.gather(table_i64, 1, clamped_blk)
valid = mask & in_bounds & (blk >= 0)
return torch.where(valid, blk * block_size + intra, -1)| def mask_uncommitted_compressed_boundaries( | ||
| mapping: torch.Tensor, | ||
| positions: torch.Tensor, | ||
| *, | ||
| compress_ratio: int, | ||
| commit_tokens: int | None, | ||
| ) -> torch.Tensor: | ||
| if commit_tokens is None: | ||
| return mapping | ||
| if mapping.shape != positions.shape: | ||
| raise ValueError("compressed boundary mask expects mapping and positions to have the same shape") | ||
| if mapping.ndim != 2: | ||
| raise ValueError("compressed boundary mask expects [B, S] tensors") | ||
| if commit_tokens < 0 or commit_tokens > mapping.shape[1]: | ||
| raise ValueError(f"commit_tokens must be in [0, {mapping.shape[1]}], got {commit_tokens}") | ||
| masked = mapping.clone() | ||
| positions_i64 = positions.to(torch.int64) | ||
| for b in range(positions_i64.shape[0]): | ||
| for s in range(commit_tokens, positions_i64.shape[1]): | ||
| pos = int(positions_i64[b, s].item()) | ||
| if (pos + 1) % compress_ratio == 0: | ||
| masked[b, s] = -1 | ||
| return masked |
There was a problem hiding this comment.
The nested loops can be replaced with a vectorized mask operation using torch.arange and broadcasting, which is much more efficient.
def mask_uncommitted_compressed_boundaries(
mapping: torch.Tensor,
positions: torch.Tensor,
*,
compress_ratio: int,
commit_tokens: int | None,
) -> torch.Tensor:
if commit_tokens is None:
return mapping
if mapping.shape != positions.shape:
raise ValueError("compressed boundary mask expects mapping and positions to have the same shape")
if mapping.ndim != 2:
raise ValueError("compressed boundary mask expects [B, S] tensors")
if commit_tokens < 0 or commit_tokens > mapping.shape[1]:
raise ValueError(f"commit_tokens must be in [0, {mapping.shape[1]}], got {commit_tokens}")
masked = mapping.clone()
positions_i64 = positions.to(torch.int64)
s_indices = torch.arange(positions.shape[1], device=positions.device).unsqueeze(0)
s_mask = s_indices >= commit_tokens
cond = s_mask & ((positions_i64 + 1) % compress_ratio == 0)
masked[cond] = -1
return masked| def state_slot_mapping( | ||
| positions: torch.Tensor, | ||
| state_block_table: torch.Tensor, | ||
| *, | ||
| state_block_size: int, | ||
| ) -> torch.Tensor: | ||
| mapping = torch.full(positions.shape, -1, dtype=torch.int64) | ||
| table_i64 = state_block_table.to(torch.int64) | ||
| positions_i64 = positions.to(torch.int64) | ||
| for b in range(positions_i64.shape[0]): | ||
| for s in range(positions_i64.shape[1]): | ||
| pos = int(positions_i64[b, s].item()) | ||
| logical_blk = pos // state_block_size | ||
| intra = pos % state_block_size | ||
| if logical_blk >= table_i64.shape[1]: | ||
| continue | ||
| blk = int(table_i64[b, logical_blk].item()) | ||
| if blk >= 0: | ||
| mapping[b, s] = blk * state_block_size + intra | ||
| return mapping |
There was a problem hiding this comment.
This function can be fully vectorized using torch.gather and torch.where, avoiding the nested loops and .item() calls.
def state_slot_mapping(
positions: torch.Tensor,
state_block_table: torch.Tensor,
*,
state_block_size: int,
) -> torch.Tensor:
positions_i64 = positions.to(torch.int64)
table_i64 = state_block_table.to(torch.int64)
logical_blk = positions_i64 // state_block_size
intra = positions_i64 % state_block_size
in_bounds = logical_blk < table_i64.shape[1]
clamped_blk = torch.clamp(logical_blk, max=table_i64.shape[1] - 1)
blk = torch.gather(table_i64, 1, clamped_blk)
valid = in_bounds & (blk >= 0)
return torch.where(valid, blk * state_block_size + intra, -1)
sjduan
force-pushed
the
dsv4-decode-8k-minimal
branch
from
8ccd859 to
51e8757
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
models/deepseek/v4/decode_metadata.py (1)
43-49: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider vectorizing the metadata builders instead of per-element Python loops.
position_ids_from_starts,block_table,ori_slot_mapping,compressed_slot_mapping, andstate_slot_mappingall build results with nested Pythonforloops and per-element.item()extraction. Per the decode_fwd.py context, these helpers appear to run inside the real repeated per-token decode path (not only test fixtures), so the Python-level overhead scales withbatch * seqon every decode step. These are straightforward to vectorize withtorchbroadcasting/torch.arange/torch.gather, avoiding the per-element.item()round-trips.Example for
position_ids_from_starts:♻️ Proposed vectorized refactor
def position_ids_from_starts(starts: torch.Tensor, *, seq: int = DECODE_SEQ) -> torch.Tensor: - starts_i64 = starts.to(torch.int64) - positions = torch.empty((int(starts_i64.numel()), seq), dtype=torch.int32) - for b in range(int(starts_i64.numel())): - for s in range(seq): - positions[b, s] = int(starts_i64[b].item()) + s - return positions + starts_i64 = starts.to(torch.int64) + offsets = torch.arange(seq, dtype=torch.int64) + return (starts_i64.unsqueeze(1) + offsets.unsqueeze(0)).to(torch.int32)Also applies to: 64-79, 82-100, 103-127, 154-174
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@models/deepseek/v4/decode_metadata.py` around lines 43 - 49, The metadata builders still use nested Python loops and per-element .item() reads, which adds avoidable overhead in the decode path. Refactor position_ids_from_starts and the related helpers block_table, ori_slot_mapping, compressed_slot_mapping, and state_slot_mapping to use vectorized torch operations such as broadcasting, torch.arange, and tensor indexing/gather instead of per-element iteration. Keep the same outputs and shapes while moving the work into tensor ops so decode_fwd.py can call them efficiently.models/deepseek/v4/decode_fwd.py (1)
999-1032: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winMetadata tensors are rebuilt once per name (15× redundant work).
Each of the 15 specs'
init_valuecallsmake_forward_metadata_tensors(...)in full and then indexes a single[name]. Materializing all specs therefore constructs the entire block-table/slot-mapping/position set 15 times. Compute the dict once and share it while preserving laziness.♻️ Compute the tensor dict once
def _make_forward_metadata_specs(base_specs, start_pos=None, commit_tokens=1): from golden import TensorSpec metadata_names = [ ... ] + _cache = {} + def _get(name): + if not _cache: + _cache.update( + make_forward_metadata_tensors( + base_specs, start_pos=start_pos, commit_tokens=commit_tokens + ) + ) + return _cache[name] + return { name: TensorSpec( name, list(base_specs[name].shape), base_specs[name].dtype, - init_value=lambda name=name: make_forward_metadata_tensors( - base_specs, - start_pos=start_pos, - commit_tokens=commit_tokens, - )[name], + init_value=lambda name=name: _get(name), ) for name in metadata_names }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@models/deepseek/v4/decode_fwd.py` around lines 999 - 1032, The metadata tensor specs in _make_forward_metadata_specs are doing redundant work by calling make_forward_metadata_tensors once per TensorSpec via each init_value. Fix this by computing the metadata tensor dict a single time and reusing it across all specs, while still preserving lazy initialization semantics. Use the existing _make_forward_metadata_specs and make_forward_metadata_tensors symbols to locate the change, and make sure each TensorSpec pulls from the shared dict instead of rebuilding it.
🤖 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.
Nitpick comments:
In `@models/deepseek/v4/decode_fwd.py`:
- Around line 999-1032: The metadata tensor specs in
_make_forward_metadata_specs are doing redundant work by calling
make_forward_metadata_tensors once per TensorSpec via each init_value. Fix this
by computing the metadata tensor dict a single time and reusing it across all
specs, while still preserving lazy initialization semantics. Use the existing
_make_forward_metadata_specs and make_forward_metadata_tensors symbols to locate
the change, and make sure each TensorSpec pulls from the shared dict instead of
rebuilding it.
In `@models/deepseek/v4/decode_metadata.py`:
- Around line 43-49: The metadata builders still use nested Python loops and
per-element .item() reads, which adds avoidable overhead in the decode path.
Refactor position_ids_from_starts and the related helpers block_table,
ori_slot_mapping, compressed_slot_mapping, and state_slot_mapping to use
vectorized torch operations such as broadcasting, torch.arange, and tensor
indexing/gather instead of per-element iteration. Keep the same outputs and
shapes while moving the work into tensor ops so decode_fwd.py can call them
efficiently.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 40ef1fd0-d2a3-4f43-8d3e-ec5762ee8296
📒 Files selected for processing (10)
models/deepseek/v4/config.pymodels/deepseek/v4/decode_attention_csa.pymodels/deepseek/v4/decode_attention_hca.pymodels/deepseek/v4/decode_attention_swa.pymodels/deepseek/v4/decode_compressor_ratio128.pymodels/deepseek/v4/decode_compressor_ratio4.pymodels/deepseek/v4/decode_fwd.pymodels/deepseek/v4/decode_indexer.pymodels/deepseek/v4/decode_indexer_compressor.pymodels/deepseek/v4/decode_metadata.py
sjduan
force-pushed
the
dsv4-decode-8k-minimal
branch
from
51e8757 to
3bec99f
Compare
|
Rebased onto One small follow-up fix was included after rebase: Validation after rebase:
Note: the older PyPTO remote machine currently fails latest |
|
Follow-up for the The previous workaround made I updated the branch to keep the public Validated on 2号机 container:
|
sjduan
force-pushed
the
dsv4-decode-8k-minimal
branch
from
3bec99f to
63bbee1
Compare
|
Updated The previous mixed style ( Validated on 2号机 container:
|
sjduan
force-pushed
the
dsv4-decode-8k-minimal
branch
from
63bbee1 to
dd06c94
Compare
2 participants
Summary
Scope Notes
Tests
Local:
Remote a2a3, EP-independent:
Remote a2a3, EP2:
Remote a2a3, EP8 smoke: