Skip to content

dsa: remove block_table_convert_triton in dsa #658

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
junhaha666 wants to merge 4 commits into from
Closed

dsa: remove block_table_convert_triton in dsa #658

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9b5e76d
dsa: remove block_table_convert_triton in dsa
junhaha666 Apr 28, 2026
c50a7e5
Merge branch 'main' into jun/dsa_fix
junhaha666 Apr 30, 2026
ed1bec0
clear code
junhaha666 Apr 30, 2026
58be241
Merge branch 'main' into jun/dsa_fix
valarLip May 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 10 additions & 4 deletions atom/model_ops/attention_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,7 @@ def _forward_prefill_mla(
attn_metadata.block_tables,
attn_metadata.cu_seqlens_k,
NUM_TOPK_TOKENS=self.topk_indices_buffer.shape[1],
PAGE_SIZE=get_current_atom_config().kv_cache_block_size,
)
paged_cu_seqlens_q = attn_metadata.sparse_cu_seqlens_q
paged_kv_indptr = attn_metadata.sparse_kv_indptr
Expand Down Expand Up @@ -916,7 +917,7 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel(
out_kv_indices, # int32
# shapes (compile-time where possible)
NUM_TOPK_TOKENS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
PAGE_SIZE: tl.constexpr,
BLOCK_N: tl.constexpr, # tile width along columns
# strides (in elements)
ti_stride0: tl.int64, # topk_indices stride 0
Expand All @@ -941,14 +942,19 @@ def _convert_req_index_to_global_index_dsa_prefill_kernel(
) # int32
pre_seqlens_q = tl.load(cu_seqlens_q + req_id)

seq_token_idx = indice - pre_seqlens_q
block_id = seq_token_idx // PAGE_SIZE
inblock_offset = seq_token_idx % PAGE_SIZE

# Guard block_table access
store_mask = (col_id < kv_len) & (col_id < NUM_TOPK_TOKENS)
valid_mask = store_mask & (indice >= 0)
out_val = tl.load(
block_table + req_id * bt_stride0 + (indice - pre_seqlens_q) * bt_stride1,
physical_block = tl.load(
block_table + req_id * bt_stride0 + block_id * bt_stride1,
mask=valid_mask,
other=-1,
)
out_val = tl.where(valid_mask, physical_block * PAGE_SIZE + inblock_offset, -1)

# Store results
out_ptr_ij = out_kv_indices + kv_start + col_id
Expand All @@ -967,7 +973,7 @@ def triton_convert_req_index_to_global_index_dsa_prefill(
block_table: torch.Tensor, # int32 [num_req, max_num_blocks_per_req]
cu_seqlens_q: torch.Tensor, # int32 [num_tokens + 1]
# dsa_kv_indices: torch.Tensor, # int32 [total_kv_seqlen] -->>> output for this kernel
PAGE_SIZE: int = 1, # page_block_size = 1 for now
PAGE_SIZE: int = 1,
NUM_TOPK_TOKENS: int = 2048,
BLOCK_N: int = 1024, # tile width along columns
):
Expand Down
26 changes: 13 additions & 13 deletions atom/model_ops/attentions/aiter_mla.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -381,14 +381,14 @@ def prepare_prefill(self, batch: ScheduledBatch):
if attn_metadata.block_tables is None:
self.prepare_block_tables(batch)
attn_metadata.block_tables = var["block_tables"].copy_to_gpu(bs)
Comment on lines 600 to 602

Copilot AI Apr 30, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In sparse prefill, CommonAttentionBuilder.prepare_prefill() will pass block_tables_converted when has_cached=True, and AttentionMetaData.__init__ overrides attn_metadata.block_tables with that converted (per-token) table. After this PR, downstream DSA code paths (e.g. triton_convert_req_index_to_global_index_dsa_prefill using PAGE_SIZE) assume attn_metadata.block_tables contains raw per-page physical block IDs; leaving the converted table here will miscompute indices (effectively applying the page math twice). Consider always resetting attn_metadata.block_tables to the raw var["block_tables"] for the sparse MLA/DSA path (even when it’s already set), and ensure block_tables_converted does not override it for this backend.

Suggested change
if attn_metadata.block_tables is None:
self.prepare_block_tables(batch)
attn_metadata.block_tables = var["block_tables"].copy_to_gpu(bs)
# Sparse MLA/DSA kernels consume raw per-page physical block IDs.
# CommonAttentionBuilder.prepare_prefill() may populate
# attn_metadata.block_tables with a converted per-token table when
# has_cached=True, so always overwrite it here with the raw backend
# block table for this path.
self.prepare_block_tables(batch)
attn_metadata.block_tables = var["block_tables"].copy_to_gpu(bs)

Copilot uses AI. Check for mistakes.
if self.block_ratio > 1:
block_table_convert_triton(
var["block_tables"].gpu[:bs],
var["block_tables_converted"].gpu[:bs],
var["context_lens"].gpu[:bs],
self.block_ratio,
)
attn_metadata.block_tables = var["block_tables_converted"].gpu[:bs]
# if self.block_ratio > 1:
# block_table_convert_triton(
Comment thread
junhaha666 marked this conversation as resolved.
Outdated
# var["block_tables"].gpu[:bs],
# var["block_tables_converted"].gpu[:bs],
# var["context_lens"].gpu[:bs],
# self.block_ratio,
# )
# attn_metadata.block_tables = var["block_tables_converted"].gpu[:bs]
Comment thread
junhaha666 marked this conversation as resolved.
Outdated
counts = var["cu_seqlens_q"].np[1 : bs + 1] - var["cu_seqlens_q"].np[:bs]
if attn_metadata.has_cached:
# Full context (cached + new): use cu_seqlens_k for indexer
Expand Down Expand Up @@ -787,11 +787,11 @@ def build_for_cudagraph_capture(self, bs: int) -> AttentionMetaData:
kv_indices=var["kv_indices"].gpu,
kv_last_page_lens=var["kv_last_page_lens"].gpu[:bs],
sparse_kv_indptr=sparse_kv_indptr,
block_tables_converted=(
var["block_tables_converted"].gpu[:bs]
if "block_tables_converted" in var
else None
),
# block_tables_converted=(
# var["block_tables_converted"].gpu[:bs]
# if "block_tables_converted" in var
# else None
# ),
Comment thread
junhaha666 marked this conversation as resolved.
Outdated
**ctx_mla_ps,
)
attn_matadata.dtype_q = self.dtype_q
Expand Down
8 changes: 7 additions & 1 deletion atom/models/deepseek_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1012,12 +1012,15 @@ def sparse_attn_indexer(
# dummy runner
return weights
num_decode_tokens = context.batch_size if not context.is_prefill else 0
ori_block_size = get_current_atom_config().kv_cache_block_size

@valarLip valarLip Apr 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

rename to runner_block_size?

kv_cache = kv_cache.view(-1, ori_block_size, kv_cache.shape[-1])
indexer_k_quant_and_cache(
k,
kv_cache,
slot_mapping,
quant_block_size,
scale_fmt,
preshuffle=True,
)
if context.is_prefill:
if attn_metadata.max_seqlen_k <= topk_indices_buffer.shape[1]:
Expand Down Expand Up @@ -1049,6 +1052,7 @@ def sparse_attn_indexer(
if prefill_metadata.has_cached
else prefill_metadata.cu_seqlens_q
),
preshuffle=True,
)
cu_seqlen_ks = prefill_metadata.cu_seqlen_ks
cu_seqlen_ke = prefill_metadata.cu_seqlen_ke
Expand Down Expand Up @@ -1100,6 +1104,8 @@ def sparse_attn_indexer(
decode_metadata.context_lens,
attn_metadata.block_tables,
max_model_len,
KVBlockSize=kv_cache.shape[1],

@valarLip valarLip Apr 30, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

use runner_block_size instead of shape.. to avoid access torch stuff

Preshuffle=True,
)
num_rows = logits.shape[0]
assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
Expand Down Expand Up @@ -1192,7 +1198,7 @@ def __init__(
self.weights_proj = ReplicatedLinear(
hidden_size,
self.n_head,
quant_config=quant_config,
quant_config=None,
Comment thread
junhaha666 marked this conversation as resolved.
prefix=f"{prefix}.weights_proj",
)
self.softmax_scale = self.head_dim**-0.5
Expand Down
Loading