From 8a7c8a72d64bce8731c40c45992ea349f966fad4 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:05:04 -0700 Subject: [PATCH] Fix NaN logits from deterministic Triton extend on the unified memory pool (#33517) Co-authored-by: Claude Fable 5 --- .../srt/layers/attention/triton_backend.py | 9 ++ .../test_unified_memory_deterministic.py | 118 ++++++++++++++++++ 2 files changed, 127 insertions(+) create mode 100644 test/registered/attention/test_unified_memory_deterministic.py diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 301c45a48..54a6c2885 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -1678,7 +1678,16 @@ class TritonAttnBackend(AttentionBackend): and isinstance(pool, SWAKVPool) and pool.layers_mapping[layer.layer_id][1] ): + # Consumes VIRTUAL ids, so it must see out_cache_loc untranslated. extend_kv_indices = pool.translate_loc_from_full_to_swa(extend_kv_indices) + elif self.forward_metadata.out_cache_loc_full_physical is not None: + # Unified pool: this kernel reads the extend half OUT OF THE POOL (the + # 2-stage path takes it from the k/v arguments), so it needs the same + # translated loc the KV write uses -- otherwise the prefix is read at + # physical ids and the extend tokens at virtual ones. Reuse the + # per-forward translation rather than re-translating: this runs once + # per layer. + extend_kv_indices = self.forward_metadata.out_cache_loc_full_physical # Handle cases where extend_seq_lens or extend_start_loc might not be set # In speculative decoding, we can infer these from spec_info or compute them diff --git a/test/registered/attention/test_unified_memory_deterministic.py b/test/registered/attention/test_unified_memory_deterministic.py new file mode 100644 index 000000000..ea7d0e111 --- /dev/null +++ b/test/registered/attention/test_unified_memory_deterministic.py @@ -0,0 +1,118 @@ +"""Deterministic inference on the unified memory pool (virtual-id KV). + +`--enable-deterministic-inference` routes Triton extend through +`TritonAttnBackend._forward_extend_unified`, a 1-stage kernel that reads BOTH +the prefix and the extend half of the KV out of the pool -- unlike the default +2-stage path, which takes the extend half from its `k`/`v` arguments. On the +unified memory pool `forward_batch.out_cache_loc` holds VIRTUAL ids, so feeding +it to that kernel untranslated read the prefix at physical ids and the extend +tokens at virtual ones. The mismatch produced garbage logits, surfacing as +`NaN detected! sampler: next_token_logits` once CI arms +`SGLANG_ENABLE_ASYNC_ASSERT` and as silent corruption without it. + +Needs all three of unified memory + Triton attention + deterministic inference; +any two are clean, which is why no existing test covered it. +""" + +import unittest + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=300, stage="base-b", runner_config="1-gpu-large") + +KIMI_LINEAR_MODEL = "yujiepan/kimi-linear-tiny-random" +SERVER_ENV = {"SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM": "0"} +# Triton is pinned on purpose here: the defect lives in the Triton deterministic +# extend kernel, so resolving to another backend would not exercise it. +BASE_ARGS = [ + "--trust-remote-code", + "--skip-tokenizer-init", + "--random-seed", + "1", + "--enable-deterministic-inference", + "--attention-backend", + "triton", + "--linear-attn-backend", + "triton", + "--mamba-backend", + "triton", + "--max-mamba-cache-size", + "32", + "--max-total-tokens", + "4096", + "--cuda-graph-backend-decode", + "disabled", + "--cuda-graph-backend-prefill", + "disabled", +] + + +class TestUnifiedMemoryDeterministicParity(CustomTestCase): + """Unified pool must match the static pool token-for-token. A multi-token + prompt is required: the defect only shows once an extend batch writes tokens + that the same forward then reads back through `out_cache_loc`.""" + + @classmethod + def setUpClass(cls): + cls.model = KIMI_LINEAR_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + + @classmethod + def _generate(cls): + response = requests.post( + cls.base_url + "/generate", + json={ + "input_ids": [1] + [100 + i % 1000 for i in range(256)], + "sampling_params": { + "temperature": 0, + "max_new_tokens": 4, + "ignore_eos": True, + }, + "return_logprob": True, + }, + timeout=120, + ) + response.raise_for_status() + return response.json()["meta_info"]["output_token_logprobs"] + + def _run(self, extra_args): + process = popen_launch_server( + self.model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=BASE_ARGS + extra_args, + env=SERVER_ENV, + ) + try: + return self._generate() + finally: + kill_process_tree(process.pid, wait_timeout=60) + + def test_matches_static_pool(self): + unified = self._run(["--enable-unified-memory"]) + static = self._run([]) + + self.assertEqual(len(unified), 4) + for logprob, _token_id, _text in unified: + self.assertEqual(logprob, logprob, "logprob is NaN") + + self.assertEqual( + [token_id for _lp, token_id, _t in unified], + [token_id for _lp, token_id, _t in static], + "unified pool diverged from the static pool", + ) + for (unified_lp, _uid, _ut), (static_lp, _sid, _st) in zip(unified, static): + self.assertAlmostEqual(unified_lp, static_lp, delta=0.05) + + +if __name__ == "__main__": + unittest.main()