From be80cac93f936e96d353e1c0c9952438bf88f289 Mon Sep 17 00:00:00 2001 From: Chunan Zeng Date: Sat, 24 Oct 2026 15:37:19 +0800 Subject: [PATCH] [Fix] attention: Make FlashInfer prefill bit-exact on radix-cache hits With --enable-deterministic-inference, FlashInfer prefill was bit-exact only when the prompt was uncached. Once the radix cache turned a prompt into a prefix hit, only the tail (q_len = seq_len - prefix_len) went to BatchPrefillWithPagedKVCacheWrapper. That changed avg_packed_qo_len and made FA2DetermineCtaTileQ pick a different cta_tile_q bucket than the full prefill; different tile sizes split the reduction differently, so cached and uncached runs diverged even at temperature 0. Under deterministic inference, pad the front of each paged-prefill q block with zero rows up to deterministic_prefill_min_q(num_qo_heads, num_kv_heads, head_dim), a per-model length that keeps avg_packed_qo_len above every bucket threshold. FlashInfer aligns the q block to the tail of the KV segment, so real rows keep their causal positions; padded rows are discarded by gathering the output back to the real rows. Padding is disabled for ragged prefill, sliding-window KV pools, custom KV indices/masks, multi-item scoring, piecewise CUDA graph capture and cuda-graph prefill wrappers: the full prefill CUDA graph plans against a static q buffer sized at capture and picks its tile from the captured max packed qo length, so raising qo_indptr there broke the q.shape[0] == qo_indptr[-1] invariant at capture time. Also: - allow flashinfer in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND and split the absorbed-MLA whitelist into MLA_SUPPORTED_... so flashinfer stays rejected for MLA models - make the radix_cache mode in test_deterministic fail when the cache did not actually hit - add TestFlashinferDeterministicRadixCache to CUDA CI; it disables the prefill CUDA graph because the breakable graph adds its own non-determinism - register the padding unit test in base-b and add a kernel-level test in base-b-kernel-unit proving a padded short extend matches the full prefill bit for bit, for single and mixed batches, and that a cuda-graph wrapper keeps the real qo_indptr - document radix cache as supported for FlashInfer Hybrid GDN models (e.g. Qwen3.5) still show cached/uncached drift for q_len <= 32 with every backend; that is a pre-existing property of the hybrid path, not of this change. Co-Authored-By: Claude Fable 5.1 Co-Authored-By: Claude Code --- .../deterministic_inference.mdx | 2 +- .../sglang/srt/arg_groups/attention_hook.py | 5 +- python/sglang/srt/arg_groups/choices.py | 10 +- .../layers/attention/flashinfer_backend.py | 107 +++++++- .../attention/flashinfer_deterministic.py | 43 ++++ python/sglang/srt/server_args.py | 1 + python/sglang/test/test_deterministic.py | 6 +- .../attention/test_deterministic.py | 23 +- .../test_flashinfer_deterministic_prefill.py | 236 ++++++++++++++++++ .../test_flashinfer_deterministic.py | 47 ++++ test/registered/unit/test_model_overrides.py | 9 + 11 files changed, 481 insertions(+), 8 deletions(-) create mode 100644 python/sglang/srt/layers/attention/flashinfer_deterministic.py create mode 100644 test/registered/kernels/ops/attention/test_flashinfer_deterministic_prefill.py create mode 100644 test/registered/unit/layers/attention/test_flashinfer_deterministic.py diff --git a/docs/docs/advanced_features/deterministic_inference.mdx b/docs/docs/advanced_features/deterministic_inference.mdx index 9c67328105d2..b2403fba8ab3 100644 --- a/docs/docs/advanced_features/deterministic_inference.mdx +++ b/docs/docs/advanced_features/deterministic_inference.mdx @@ -49,7 +49,7 @@ The following table shows feature compatibility for deterministic inference acro **FlashInfer** ✅ Yes ✅ Yes - ❌ No + ✅ Yes ✅ Yes diff --git a/python/sglang/srt/arg_groups/attention_hook.py b/python/sglang/srt/arg_groups/attention_hook.py index 9fadd2aee18e..f4b736c7f268 100644 --- a/python/sglang/srt/arg_groups/attention_hook.py +++ b/python/sglang/srt/arg_groups/attention_hook.py @@ -583,6 +583,7 @@ def handle_multi_item_scoring(server_args: Any): def handle_deterministic_inference(server_args: Any): from sglang.srt.server_args import ( + MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND, RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND, ) @@ -645,9 +646,9 @@ def handle_deterministic_inference(server_args: Any): attention_backend = resolved_view(server_args).attention_backend if is_deepseek_model: - if attention_backend not in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND: + if attention_backend not in MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND: raise ValueError( - f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with absorbed-MLA models. But you're using {attention_backend}." + f"Currently only {MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with absorbed-MLA models. But you're using {attention_backend}." ) if attention_backend == "fa4" and not get_platform().is_sm100_or_sm110: raise ValueError( diff --git a/python/sglang/srt/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index b79f02ee62fa..ee33dfbbe915 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/python/sglang/srt/arg_groups/choices.py @@ -246,7 +246,15 @@ DETERMINISTIC_ATTENTION_BACKEND_CHOICES.extend ) -RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["ascend", "fa3", "fa4", "triton"] +MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["ascend", "fa3", "fa4", "triton"] + +RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = [ + "ascend", + "flashinfer", + "fa3", + "fa4", + "triton", +] add_radix_supported_deterministic_attention_backend_choices = ( RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND.extend ) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index da879da053a7..2a52b848d7dc 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -35,6 +35,10 @@ resolve_window_id_space, resolve_window_kv_pool, ) +from sglang.srt.layers.attention.flashinfer_deterministic import ( + build_prefill_pad_indices, + deterministic_prefill_min_q, +) from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( KVCacheAttentionAccessKind, ) @@ -172,6 +176,8 @@ class PrefillMetadata: extend_no_prefix: bool multi_item_params: Optional[MultiItemScoringParams] = None swa_out_cache_loc: Optional[torch.Tensor] = None + prefill_q_pad_gather_idx: Optional[torch.Tensor] = None + prefill_q_pad_out_idx: Optional[torch.Tensor] = None # Reuse this workspace buffer across all flashinfer wrappers @@ -1094,6 +1100,8 @@ def init_forward_metadata(self, forward_batch: ForwardBatch): extend_no_prefix, multi_item_params, swa_out_cache_loc=swa_out_cache_loc, + prefill_q_pad_gather_idx=self.indices_updater_prefill.prefill_q_pad_gather_idx, + prefill_q_pad_out_idx=self.indices_updater_prefill.prefill_q_pad_out_idx, ) def init_cuda_graph_state( @@ -1431,8 +1439,14 @@ def forward_extend( not layer.is_cross_attention and layer.attn_type != AttentionType.ENCODER_ONLY ) + q_3d = q.view(-1, layer.tp_q_head_num, layer.head_dim) + pad_gather_idx = self.forward_metadata.prefill_q_pad_gather_idx + if pad_gather_idx is not None: + q_3d = torch.cat( + [q_3d, q_3d.new_zeros((1, *q_3d.shape[1:]))], dim=0 + ).index_select(0, pad_gather_idx) o = prefill_wrapper_paged.forward( - q.view(-1, layer.tp_q_head_num, layer.head_dim), + q_3d, kv_cache, causal=causal, sm_scale=layer.scaling, @@ -1442,6 +1456,8 @@ def forward_extend( k_scale=layer.k_scale_float, v_scale=layer.v_scale_float, ) + if pad_gather_idx is not None: + o = o.index_select(0, self.forward_metadata.prefill_q_pad_out_idx) else: # If `k`/`v` are not explicitly provided, fall back to the KV cache stored in # `self.token_to_kv_pool` for this layer. This enables attention over @@ -1925,6 +1941,21 @@ def __init__(self, model_runner: ModelRunner, attn_backend: FlashInferAttnBacken self._swa_space = attn_backend._swa_space self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged + self.prefill_cta_min_q = ( + deterministic_prefill_min_q( + self.num_qo_heads, self.num_kv_heads, self.head_dim + ) + if attn_backend.enable_deterministic + else 1 + ) + self.prefill_q_pad_enabled = ( + attn_backend.enable_deterministic + and attn_backend.dispatch_reason is None + and not attn_backend.enable_mis + ) + self.prefill_q_pad_gather_idx: Optional[torch.Tensor] = None + self.prefill_q_pad_out_idx: Optional[torch.Tensor] = None + # Dispatch the update function if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW: self.update = self.update_sliding_window @@ -2003,6 +2034,7 @@ def update_single_wrapper( fixed_split_size=fixed_split_size, multi_item_params=multi_item_params, seq_lens_cpu=seq_lens_cpu, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, custom_kv_indices=custom_kv_indices, plan=plan, ) @@ -2211,6 +2243,48 @@ def update_cross_attention( plan=plan, ) + def _can_pad_prefill_q( + self, + wrapper_paged: BatchPrefillWithPagedKVCacheWrapper, + use_ragged: bool, + use_sliding_window_kv_pool: bool, + custom_kv_indices: Optional[torch.Tensor], + cross_attention_custom_mask: Optional[torch.Tensor], + multi_item_params: Optional[MultiItemScoringParams], + ) -> bool: + if not self.prefill_q_pad_enabled or self.prefill_cta_min_q <= 1: + return False + if wrapper_paged.is_cuda_graph_enabled: + return False + if use_ragged or use_sliding_window_kv_pool: + return False + if custom_kv_indices is not None or cross_attention_custom_mask is not None: + return False + if multi_item_params is not None and multi_item_params.is_enabled(): + return False + return not is_in_tc_piecewise_cuda_graph() + + def _prefill_q_lens_cpu( + self, + q_lens: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + extend_prefix_lens_cpu: Optional[List[int]], + ) -> List[int]: + bs = q_lens.shape[0] + if ( + extend_prefix_lens_cpu is not None + and seq_lens_cpu is not None + and len(extend_prefix_lens_cpu) == bs + and len(seq_lens_cpu) == bs + ): + return [ + int(seq_len) - int(prefix_len) + for seq_len, prefix_len in zip( + seq_lens_cpu.tolist(), extend_prefix_lens_cpu + ) + ] + return q_lens.detach().to("cpu", torch.int64).tolist() + def call_begin_forward( self, wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper, @@ -2230,12 +2304,15 @@ def call_begin_forward( multi_item_params: Optional[MultiItemScoringParams] = None, cross_attention_custom_mask: Optional[torch.Tensor] = None, seq_lens_cpu: Optional[torch.Tensor] = None, + extend_prefix_lens_cpu: Optional[List[int]] = None, custom_kv_indices: Optional[torch.Tensor] = None, window: Optional[WindowRead] = None, *, plan: IterationIds, ): bs = len(seq_lens) + self.prefill_q_pad_gather_idx = None + self.prefill_q_pad_out_idx = None # Unified SWA wrapper-0: gather from the swa canonical directly -- its # entries are already swa-side physical ids, so the in-place # full->swa translate below must not run on top of them. @@ -2276,7 +2353,33 @@ def call_begin_forward( kv_start_idx=kv_start_idx, kind=read_kind, ) - qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0) + q_lens = seq_lens - prefix_lens + if self._can_pad_prefill_q( + wrapper_paged=wrapper_paged, + use_ragged=use_ragged, + use_sliding_window_kv_pool=use_sliding_window_kv_pool, + custom_kv_indices=custom_kv_indices, + cross_attention_custom_mask=cross_attention_custom_mask, + multi_item_params=multi_item_params, + ): + q_lens_cpu = self._prefill_q_lens_cpu( + q_lens, seq_lens_cpu, extend_prefix_lens_cpu + ) + gather_idx, out_idx = build_prefill_pad_indices( + q_lens_cpu, self.prefill_cta_min_q, seq_lens.device + ) + if gather_idx is not None: + self.prefill_q_pad_gather_idx = gather_idx + self.prefill_q_pad_out_idx = out_idx + q_lens = torch.tensor( + [ + max(length, self.prefill_cta_min_q) if length > 0 else 0 + for length in q_lens_cpu + ], + dtype=q_lens.dtype, + device=q_lens.device, + ) + qo_indptr[1 : bs + 1] = torch.cumsum(q_lens, dim=0) qo_indptr = qo_indptr[: bs + 1] custom_mask = cross_attention_custom_mask diff --git a/python/sglang/srt/layers/attention/flashinfer_deterministic.py b/python/sglang/srt/layers/attention/flashinfer_deterministic.py new file mode 100644 index 000000000000..1cfab1839032 --- /dev/null +++ b/python/sglang/srt/layers/attention/flashinfer_deterministic.py @@ -0,0 +1,43 @@ +from typing import List, Optional, Sequence, Tuple + +import torch + + +def deterministic_prefill_min_q( + num_qo_heads: int, num_kv_heads: int, head_dim: int +) -> int: + gqa_group = max(num_qo_heads // num_kv_heads, 1) + if head_dim >= 512: + threshold = 32 + elif head_dim >= 256: + threshold = 16 + else: + threshold = 64 + return threshold // gqa_group + 1 + + +def build_prefill_pad_indices( + real_lens: Sequence[int], min_q: int, device: torch.device +) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + padded_lens: List[int] = [max(real, min_q) if real > 0 else 0 for real in real_lens] + real_total = sum(real_lens) + padded_total = sum(padded_lens) + if padded_total == real_total: + return None, None + + gather: List[int] = [] + out: List[int] = [] + real_off = 0 + pad_off = 0 + for real, padded in zip(real_lens, padded_lens): + lead = padded - real + gather.extend([real_total] * lead) + gather.extend(range(real_off, real_off + real)) + out.extend(range(pad_off + lead, pad_off + padded)) + real_off += real + pad_off += padded + + return ( + torch.tensor(gather, dtype=torch.long, device=device), + torch.tensor(out, dtype=torch.long, device=device), + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 849cdfffe0a3..bd7807c6257e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -111,6 +111,7 @@ def _real_kv_hash_modes(): GRAMMAR_BACKEND_CHOICES, LINEAR_ATTN_KERNEL_BACKEND_CHOICES, LOAD_FORMAT_CHOICES, + MLA_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND, MOE_RUNNER_BACKEND_CHOICES, MXFP8_MOE_RUNNER_BACKEND_CHOICES, QUANTIZATION_CHOICES, diff --git a/python/sglang/test/test_deterministic.py b/python/sglang/test/test_deterministic.py index 8d02a674a894..67e451a97f85 100644 --- a/python/sglang/test/test_deterministic.py +++ b/python/sglang/test/test_deterministic.py @@ -623,8 +623,12 @@ def test_deterministic(args): cached_token_data = cached_logprobs[0] cached_logprob = cached_token_data[0] cached_token_id = cached_token_data[1] + cached_tokens = cached_response["meta_info"].get("cached_tokens", 0) - print(f"✓ Generated with cache:") + print(f"✓ Generated with cache (cached_tokens={cached_tokens}):") + if cached_tokens == 0: + print("TEST FAILED - Radix cache did not hit") + return [0] print(f" Token ID: {cached_token_id}") print(f" Logprob: {cached_logprob:.10f}") diff --git a/test/registered/attention/test_deterministic.py b/test/registered/attention/test_deterministic.py index d442b68c2dc1..38efde287711 100644 --- a/test/registered/attention/test_deterministic.py +++ b/test/registered/attention/test_deterministic.py @@ -15,16 +15,18 @@ register_cuda_ci, register_xpu_ci, ) +from sglang.test.test_deterministic import BenchArgs, test_deterministic from sglang.test.test_deterministic_utils import ( COMMON_SERVER_ARGS, TestDeterministicBase, ) from sglang.test.test_utils import ( DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, + DEFAULT_URL_FOR_TEST, is_in_amd_ci, ) -register_cuda_ci(est_time=285, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=385, stage="base-b", runner_config="1-gpu-large") register_amd_ci(est_time=278, suite="stage-b-test-1-gpu-small-amd") register_xpu_ci(est_time=207, suite="stage-b-test-1-gpu-xpu") @@ -47,6 +49,25 @@ def get_server_args(cls): return args +@unittest.skipIf(_is_xpu, "CUDA runner only") +@unittest.skipIf(is_in_amd_ci(), "Skip for AMD CI.") +class TestFlashinferDeterministicRadixCache(TestDeterministicBase): + @classmethod + def get_server_args(cls): + return list(COMMON_SERVER_ARGS) + [ + "--attention-backend", + "flashinfer", + "--disable-prefill-cuda-graph", + ] + + def test_radix_cache(self): + args = BenchArgs() + args.host, args.port = self._extract_host_and_port(DEFAULT_URL_FOR_TEST) + args.test_mode = "radix_cache" + results = test_deterministic(args) + self.assertEqual(results, [1]) + + @unittest.skipIf(_is_xpu, "CUDA runner only") @unittest.skipIf(is_in_amd_ci(), "Skip for AMD CI.") class TestFa3Deterministic(TestDeterministicBase): diff --git a/test/registered/kernels/ops/attention/test_flashinfer_deterministic_prefill.py b/test/registered/kernels/ops/attention/test_flashinfer_deterministic_prefill.py new file mode 100644 index 000000000000..329f70490684 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_flashinfer_deterministic_prefill.py @@ -0,0 +1,236 @@ +"""Padded short paged-prefill extends must match the full prefill bit for bit.""" + +import unittest +from types import SimpleNamespace + +import torch +from flashinfer.prefill import BatchPrefillWithPagedKVCacheWrapper + +from sglang.srt.layers.attention.flashinfer_backend import ( + FlashInferIndicesUpdaterPrefill, +) +from sglang.srt.layers.attention.flashinfer_deterministic import ( + deterministic_prefill_min_q, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +_HQ = 32 +_HKV = 8 +_DIM = 128 +_DTYPE = torch.bfloat16 +_MIN_Q = deterministic_prefill_min_q(_HQ, _HKV, _DIM) +_FIXED_SPLIT = 16384 +_WORKSPACE_BYTES = 1 << 30 +_DEVICE = "cuda" + + +def _indptr(lengths): + out = [0] + for length in lengths: + out.append(out[-1] + length) + return torch.tensor(out, dtype=torch.int32, device=_DEVICE) + + +def _kv_cache(kv_lens): + total = sum(kv_lens) + cache = torch.empty(total, 2, 1, _HKV, _DIM, dtype=_DTYPE, device=_DEVICE) + cache[:, 0, 0] = torch.randn(total, _HKV, _DIM, dtype=_DTYPE, device=_DEVICE) + cache[:, 1, 0] = torch.randn(total, _HKV, _DIM, dtype=_DTYPE, device=_DEVICE) + return cache + + +def _run(q_rows, qo_lens, kv_cache, kv_lens): + wrapper = BatchPrefillWithPagedKVCacheWrapper( + torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=_DEVICE), "NHD" + ) + wrapper.plan( + _indptr(qo_lens), + _indptr(kv_lens), + torch.arange(sum(kv_lens), dtype=torch.int32, device=_DEVICE), + torch.ones(len(kv_lens), dtype=torch.int32, device=_DEVICE), + _HQ, + _HKV, + _DIM, + 1, + causal=True, + fixed_split_size=_FIXED_SPLIT, + sm_scale=_DIM**-0.5, + q_data_type=_DTYPE, + kv_data_type=_DTYPE, + ) + return wrapper.run(q_rows, kv_cache) + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class TestDeterministicPrefillPadding(CustomTestCase): + def _check_short_extend(self, total_len, prefix_len): + torch.manual_seed(0) + kv_cache = _kv_cache([total_len]) + q = torch.randn(total_len, _HQ, _DIM, dtype=_DTYPE, device=_DEVICE) + full = _run(q, [total_len], kv_cache, [total_len]) + + extend = total_len - prefix_len + self.assertLess(extend, _MIN_Q) + padded_rows = torch.cat( + [ + torch.zeros(_MIN_Q - extend, _HQ, _DIM, dtype=_DTYPE, device=_DEVICE), + q[prefix_len:], + ], + dim=0, + ) + padded = _run(padded_rows, [_MIN_Q], kv_cache, [total_len]) + + self.assertTrue( + torch.equal(full[prefix_len:], padded[_MIN_Q - extend :]), + f"cached rows diverged: total={total_len} prefix={prefix_len}", + ) + + def test_single_token_extend(self): + self._check_short_extend(101, 100) + + def test_three_token_extend(self): + self._check_short_extend(98, 95) + + def test_mixed_batch_keeps_short_and_long_requests_exact(self): + torch.manual_seed(1) + kv_lens = [100, 20] + kv_cache = _kv_cache(kv_lens) + q_a = torch.randn(100, _HQ, _DIM, dtype=_DTYPE, device=_DEVICE) + q_b = torch.randn(20, _HQ, _DIM, dtype=_DTYPE, device=_DEVICE) + full = _run(torch.cat([q_a, q_b], dim=0), kv_lens, kv_cache, kv_lens) + + prefix_a, extend_a = 96, 4 + padded_rows = torch.cat( + [ + torch.zeros(_MIN_Q - extend_a, _HQ, _DIM, dtype=_DTYPE, device=_DEVICE), + q_a[prefix_a:], + q_b, + ], + dim=0, + ) + padded = _run(padded_rows, [_MIN_Q, 20], kv_cache, kv_lens) + + self.assertTrue( + torch.equal(full[prefix_a:100], padded[_MIN_Q - extend_a : _MIN_Q]), + "short request diverged in a mixed batch", + ) + self.assertTrue( + torch.equal(full[100:], padded[_MIN_Q:]), + "long request diverged when batched with a padded request", + ) + + +def _updater(bs): + def pack_read_stream(plan, **kw): + out = kw["out"] + offset = 0 + for length in kw["seq_lens"].tolist(): + out[offset : offset + length] = torch.arange( + length, dtype=torch.int32, device=_DEVICE + ) + offset += length + + updater = FlashInferIndicesUpdaterPrefill.__new__(FlashInferIndicesUpdaterPrefill) + updater.attn_backend = SimpleNamespace( + runner_ids=SimpleNamespace( + reads_are_translated=False, pack_read_stream=pack_read_stream + ), + dq_paged_kernel_lens=None, + ) + updater.num_qo_heads, updater.num_kv_heads, updater.head_dim = _HQ, _HKV, _DIM + updater.q_data_type = updater.data_type = _DTYPE + updater.kv_last_page_len = torch.ones(bs, dtype=torch.int32, device=_DEVICE) + updater._swa_kv_pool = None + updater.prefill_q_pad_enabled = True + updater.prefill_cta_min_q = _MIN_Q + return updater + + +def _paged_wrapper(bs, capacity, use_cuda_graph): + kwargs = {} + if use_cuda_graph: + kwargs = dict( + use_cuda_graph=True, + qo_indptr_buf=torch.zeros(bs + 1, dtype=torch.int32, device=_DEVICE), + paged_kv_indptr_buf=torch.zeros(bs + 1, dtype=torch.int32, device=_DEVICE), + paged_kv_indices_buf=torch.zeros( + capacity + 256, dtype=torch.int32, device=_DEVICE + ), + paged_kv_last_page_len_buf=torch.ones( + bs, dtype=torch.int32, device=_DEVICE + ), + ) + return BatchPrefillWithPagedKVCacheWrapper( + torch.empty(_WORKSPACE_BYTES, dtype=torch.uint8, device=_DEVICE), + "NHD", + backend="fa2", + **kwargs, + ) + + +def _begin_forward(updater, wrapper, total_len, prefix_len): + seq_lens = torch.tensor([total_len], dtype=torch.int32, device=_DEVICE) + prefix_lens = torch.tensor([prefix_len], dtype=torch.int32, device=_DEVICE) + kv_indptr = torch.zeros(2, dtype=torch.int32, device=_DEVICE) + qo_indptr = torch.zeros(2, dtype=torch.int32, device=_DEVICE) + updater.call_begin_forward( + None, + wrapper, + torch.zeros(1, dtype=torch.int32, device=_DEVICE), + seq_lens, + total_len, + seq_lens, + prefix_lens, + None, + kv_indptr, + qo_indptr, + False, + None, + fixed_split_size=_FIXED_SPLIT, + seq_lens_cpu=torch.tensor([total_len], dtype=torch.int32), + extend_prefix_lens_cpu=[prefix_len], + plan=None, + ) + return qo_indptr + + +@unittest.skipUnless(torch.cuda.is_available(), "CUDA required") +class TestDeterministicPrefillPaddingCudaGraph(CustomTestCase): + def test_predicate_rejects_cuda_graph_wrapper(self): + updater = _updater(1) + common = dict( + use_ragged=False, + use_sliding_window_kv_pool=False, + custom_kv_indices=None, + cross_attention_custom_mask=None, + multi_item_params=None, + ) + self.assertTrue( + updater._can_pad_prefill_q( + wrapper_paged=SimpleNamespace(is_cuda_graph_enabled=False), **common + ) + ) + self.assertFalse( + updater._can_pad_prefill_q( + wrapper_paged=SimpleNamespace(is_cuda_graph_enabled=True), **common + ) + ) + + def test_eager_wrapper_pads_short_extend(self): + updater = _updater(1) + qo_indptr = _begin_forward(updater, _paged_wrapper(1, 101, False), 101, 100) + self.assertIsNotNone(updater.prefill_q_pad_gather_idx) + self.assertEqual(qo_indptr[1].item(), _MIN_Q) + + def test_cuda_graph_wrapper_keeps_real_qo_indptr(self): + updater = _updater(1) + qo_indptr = _begin_forward(updater, _paged_wrapper(1, 101, True), 101, 100) + self.assertIsNone(updater.prefill_q_pad_gather_idx) + self.assertEqual(qo_indptr[1].item(), 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/layers/attention/test_flashinfer_deterministic.py b/test/registered/unit/layers/attention/test_flashinfer_deterministic.py new file mode 100644 index 000000000000..0ae0b5d4034b --- /dev/null +++ b/test/registered/unit/layers/attention/test_flashinfer_deterministic.py @@ -0,0 +1,47 @@ +import unittest + +from sglang.srt.layers.attention.flashinfer_deterministic import ( + build_prefill_pad_indices, + deterministic_prefill_min_q, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") + + +class TestDeterministicPrefillMinQ(CustomTestCase): + def test_head_dim_256_threshold(self): + self.assertEqual(deterministic_prefill_min_q(24, 4, 256), 3) + self.assertEqual(deterministic_prefill_min_q(16, 4, 256), 5) + + def test_head_dim_below_256_threshold(self): + self.assertEqual(deterministic_prefill_min_q(32, 8, 128), 17) + self.assertEqual(deterministic_prefill_min_q(4, 4, 64), 65) + + def test_head_dim_at_least_512_threshold(self): + self.assertEqual(deterministic_prefill_min_q(8, 1, 512), 5) + + def test_group_larger_than_threshold_needs_no_padding(self): + self.assertEqual(deterministic_prefill_min_q(32, 1, 256), 1) + + +class TestBuildPrefillPadIndices(CustomTestCase): + def test_pads_front_of_each_request(self): + gather, out = build_prefill_pad_indices([1, 3], 5, "cpu") + self.assertEqual(gather.tolist(), [4, 4, 4, 4, 0, 4, 4, 1, 2, 3]) + self.assertEqual(out.tolist(), [4, 7, 8, 9]) + + def test_no_padding_returns_none(self): + gather, out = build_prefill_pad_indices([5, 7], 5, "cpu") + self.assertIsNone(gather) + self.assertIsNone(out) + + def test_zero_length_request_is_not_padded(self): + gather, out = build_prefill_pad_indices([0, 2], 5, "cpu") + self.assertEqual(gather.tolist(), [2, 2, 2, 0, 1]) + self.assertEqual(out.tolist(), [3, 4]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index bd300394a1db..e11e4d7f8103 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1964,6 +1964,15 @@ def test_deterministic_rocm_rccl_check(self): else: self.assertEqual(_deterministic_rocm_rccl_check(view), {}) + def test_deterministic_flashinfer_keeps_radix_cache(self): + sa = self._construct( + "LlamaForCausalLM", + "llama", + enable_deterministic_inference=True, + attention_backend="flashinfer", + ) + self.assertFalse(self._resolved(sa, "disable_radix_cache")) + def test_dllm_forces_flashinfer_with_cuda_graph(self): # CUDA path: cuda graph enabled by default -> dllm forces flashinfer. # A real dllm arch: the page pass now runs regardless of the radix