[DSA] Re-enable fused top-k v2 for MTP: clamp padded-row seq_lens to >= 0 (#30378)

Co-authored-by: ziyi.xu <ziyi.xu@radixark.ai>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-07-07 13:44:01 -07:00
committed by GitHub
co-authored by ziyi.xu Claude Fable 5
parent 6875df3378
commit bbc537035a
6 changed files with 57 additions and 28 deletions
@@ -640,6 +640,12 @@ class TestDSAIndexer(CustomTestCase):
)
).contiguous()
# The fused v2 PAGED dispatch requires the per-forward plan to be
# preprocessed alongside the metadata (it asserts rather than silently
# recomputing it) -- mirror what init_forward_metadata /
# _build_forward_metadata_cuda_graph do.
from sglang.jit_kernel.dsv4.topk import plan_topk_v2
attn_metadata = DSAMetadata(
page_size=1,
cache_seqlens_int32=seq_lens_expanded.clone(),
@@ -654,6 +660,7 @@ class TestDSAIndexer(CustomTestCase):
dsa_cu_seqlens_k=dsa_cu_seqlens_k,
dsa_extend_seq_lens_list=seq_lens_expanded.cpu().tolist(),
dsa_seqlens_expanded=seq_lens_expanded,
topk_v2_plan=plan_topk_v2(seq_lens_expanded),
topk_indices_offset=(
topk_indices_offset
if topk_transform_method == TopkTransformMethod.RAGGED