[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:
co-authored by
ziyi.xu
Claude Fable 5
parent
6875df3378
commit
bbc537035a
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user