Skip to content

Remove host syncs from compute_dacs_segsum_triton_varlen #969

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

Open
jgcb00 wants to merge 2 commits into
base: main
Choose a base branch
from
+15 −18
Open

Remove host syncs from compute_dacs_segsum_triton_varlen #969

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3e4a6b2
Remove host syncs from `compute_dacs_segsum_triton_varlen`
jgcb00 Jun 11, 2026
528b405
Merge pull request #1 from jgcb00/jgcb00-patch
jgcb00 Jun 11, 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
33 changes: 15 additions & 18 deletions mamba_ssm/ops/triton/mamba3/mamba3_mimo_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -648,28 +648,25 @@ def compute_dacs_segsum_triton_varlen(
da_cs_rev = torch.empty_like(da)
segsum = torch.zeros(B, H, nchunks, chunk_size, chunk_size, device=da.device, dtype=da.dtype)

seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
chunks_per_seq = (seq_lens // chunk_size) + 1

# Build mapping tensors: both have length nchunks_global.
# Inactive (padding) slots are given a sentinel local-chunk index that
# places chunk_start >= seq_end, making the kernel mask all-False.
state_seq_mapping = torch.zeros(nchunks, dtype=torch.int32, device=da.device)
state_chunk_in_seq = torch.zeros(nchunks, dtype=torch.int32, device=da.device)
default_seq_len = int(seq_lens[0].item()) if num_sequences > 0 else 0
default_inactive_local_chunk = (default_seq_len + chunk_size - 1) // chunk_size
state_chunk_in_seq.fill_(default_inactive_local_chunk)

for i in range(num_sequences):
start = int(cu_seqlens[i].item())
n = int(chunks_per_seq[i].item())
chunk_start = (start // chunk_size) + i
chunk_end = chunk_start + n
assert chunk_end <= nchunks, (
f"Chunk mapping overflow for seq {i}: [{chunk_start}, {chunk_end}) vs nchunks={nchunks}"
)
state_seq_mapping[chunk_start:chunk_end] = i
state_chunk_in_seq[chunk_start:chunk_end] = torch.arange(n, dtype=torch.int32, device=da.device)
range_starts = cu_seqlens[:-1] // chunk_size + seq_idx
range_ends = range_starts + seq_lens // chunk_size + 1
g = torch.arange(nchunks, dtype=cu_seqlens.dtype, device=da.device)
owner = torch.searchsorted(range_starts, g, right=True) - 1
active = g < range_ends[owner]
# Sentinel for inactive slots: ceil(len_0 / C), matching the original.
default_inactive_local_chunk = (
(seq_lens[0] + chunk_size - 1) // chunk_size
)
state_seq_mapping = torch.where(
active, owner, torch.zeros_like(owner)
).to(torch.int32)
state_chunk_in_seq = torch.where(
active, g - range_starts[owner], default_inactive_local_chunk
).to(torch.int32)

grid = (B, H, nchunks)
dacs_segsum_kernel_varlen[grid](
Expand Down