From a11dba1a01ec553b01c76858ae53a2e04a056fca Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 3 Sep 2026 03:37:12 -0700 Subject: [PATCH] [Feature] Unified memory: support decode context parallelism for the trtllm_mla family (#37693) --- .../kernels/ops/attention/dcp_kernels.py | 20 +- python/sglang/srt/arg_groups/kv_cache_hook.py | 10 +- .../layers/attention/cutedsl_mla_backend.py | 29 +-- .../attention/tokenspeed_mla_backend.py | 2 +- .../layers/attention/trtllm_mla_backend.py | 23 +- .../srt/mem_cache/kv_index_translator.py | 16 ++ .../test_trtllm_mla_family_dcp_metadata.py | 219 ++++++++++++++++++ ...imi_linear_unified_memory_dcp_blackwell.py | 87 +++++++ 8 files changed, 378 insertions(+), 28 deletions(-) create mode 100644 test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py diff --git a/python/sglang/kernels/ops/attention/dcp_kernels.py b/python/sglang/kernels/ops/attention/dcp_kernels.py index eaabcd1b7..1a06d660d 100644 --- a/python/sglang/kernels/ops/attention/dcp_kernels.py +++ b/python/sglang/kernels/ops/attention/dcp_kernels.py @@ -82,13 +82,23 @@ def create_mla_kv_page_table_for_dcp( req_pool_indices_ptr, local_seq_lens_ptr, block_kv_indices_ptr, + v2p_ptr, # in: [num_pages + 1] int64 -- virtual->physical page table req_to_token_stride: tl.constexpr, block_table_stride: tl.constexpr, + mult, # runtime: kernel_page_multiplier of the target sub-pool PHYSICAL_PAGE_SIZE: tl.constexpr, DCP_SIZE: tl.constexpr, DCP_RANK: tl.constexpr, PAGES_PER_BLOCK: tl.constexpr, + HAS_V2P: tl.constexpr, ): + """This rank's cyclic slice of each request, as a page table. + + ``HAS_V2P`` picks the id space the emitted page number is in: the + DCP-collapsed page IS physical on a static pool, and still VIRTUAL under + the unified memory pool, where it takes one more gather through ``v2p_ptr`` + and a ``mult`` scale to reach the per-layer views. + """ req = tl.program_id(0) page_block = tl.program_id(1) page_offsets = page_block * PAGES_PER_BLOCK + tl.arange(0, PAGES_PER_BLOCK) @@ -102,10 +112,16 @@ def create_mla_kv_page_table_for_dcp( mask=mask, other=0, ) - physical_pages = virtual_locs // DCP_SIZE // PHYSICAL_PAGE_SIZE + pages = virtual_locs // DCP_SIZE // PHYSICAL_PAGE_SIZE + if HAS_V2P: + # A `-1` in req_to_token and a freed (`-1`) v2p row both clamp to entry + # 0, the reserved padding page. + pages = tl.where(virtual_locs < 0, 0, pages) + physical = tl.load(v2p_ptr + pages, mask=mask, other=0) + pages = tl.maximum(physical * mult, 0) tl.store( block_kv_indices_ptr + req * block_table_stride + page_offsets, - physical_pages, + pages.to(tl.int32), mask=mask, ) diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 6567f0a8e..bb228e6fc 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -305,12 +305,10 @@ def _validate_unified_memory_dcp(server_args: Any) -> None: "transfer, where translate_kv_indices_for_transfer would abort a " "server that had already booted." ) - # trtllm_mla (and its cutedsl_mla / tokenspeed_mla subclasses) build the - # MLA block table straight from req_to_token with - # create_flashmla_kv_indices_triton, whose v2p gather assumes UNWIDENED - # page ids; the DCP variant (create_mla_kv_page_table_for_dcp) has no v2p - # gather at all. Wire one of them through the other to add those here. - dcp_allowed = {"flashinfer"} + # The trtllm_mla family builds its DCP block table through the pool's v2p + # gather (create_mla_kv_page_table_for_dcp), so it speaks the same + # two-stage contract as flashinfer. + dcp_allowed = {"flashinfer", "trtllm_mla", "cutedsl_mla", "tokenspeed_mla"} backends = set(attention_backends_of(resolved_view(server_args))) backends.discard(None) assert backends <= dcp_allowed, ( diff --git a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py index 4c26f5f19..f223c5534 100644 --- a/python/sglang/srt/layers/attention/cutedsl_mla_backend.py +++ b/python/sglang/srt/layers/attention/cutedsl_mla_backend.py @@ -174,22 +174,17 @@ class CuteDslMLABackend(TRTLLMMLABackend): if self.data_type == torch.float8_e4m3fn: assert q_rope is not None and k_rope is not None if cos_sin_cache is None: - if ( - save_kv_cache - and self._fused_set_kv_concat_q_fp8 - and not self.kv_index_translator.is_translating - ): - # Static pool: out_cache_loc is already the physical loc. - # Fused: bf16->fp8 quantize + KV scatter + q concat in one - # launch; None when not covered. - query = self._set_kv_and_concat_q_fp8_fused( - layer=layer, - loc=forward_batch.out_cache_loc, - q=q, - q_rope=q_rope, - k=k, - k_rope=k_rope, - ) + if save_kv_cache and self._fused_set_kv_concat_q_fp8: + loc = self._resolve_fused_write_loc(forward_batch) + if loc is not None: + query = self._set_kv_and_concat_q_fp8_fused( + layer=layer, + loc=loc, + q=q, + q_rope=q_rope, + k=k, + k_rope=k_rope, + ) if query is None: q, k, k_rope = mla_quantize_without_rope_for_fp8( q, q_rope, k.squeeze(1), k_rope.squeeze(1) @@ -211,7 +206,7 @@ class CuteDslMLABackend(TRTLLMMLABackend): if query is None and save_kv_cache: assert k is not None and k_rope is not None self.token_to_kv_pool.set_mla_kv_buffer( - layer, forward_batch.out_cache_loc, k, k_rope + layer, self._kv_write_loc(forward_batch), k, k_rope ) if query is not None: diff --git a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py index 353cd6ffd..aceb69b19 100644 --- a/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py +++ b/python/sglang/srt/layers/attention/tokenspeed_mla_backend.py @@ -389,7 +389,7 @@ class TokenspeedMLABackend(TRTLLMMLABackend): if save_kv_cache: self.token_to_kv_pool.set_mla_kv_buffer( - layer, forward_batch.out_cache_loc, k, k_rope + layer, self._kv_write_loc(forward_batch), k, k_rope ) query = q.view(-1, layer.tp_q_head_num, layer.head_dim) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index efb7bfe85..8bf2eff96 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -377,6 +377,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ) -> None: parallel = get_parallel() pages_per_block = get_num_page_per_block_flashmla(self.page_size) + # None on a static pool, whose collapsed page is already physical. + v2p = self.kv_index_translator.full_v2p_table create_mla_kv_page_table_for_dcp[ ( block_kv_indices.shape[0], @@ -389,12 +391,15 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): req_pool_indices, local_seq_lens, block_kv_indices, + v2p, self.req_to_token.stride(0), block_kv_indices.stride(0), + self.kv_index_translator.full_page_multiplier, PHYSICAL_PAGE_SIZE=self.page_size, DCP_SIZE=parallel.dcp_size, DCP_RANK=parallel.dcp_rank, PAGES_PER_BLOCK=pages_per_block, + HAS_V2P=v2p is not None, ) def _create_block_kv_indices( @@ -777,6 +782,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): else: self._decode_kernel_loc = None + def _kv_write_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: + """The loc an unfused KV scatter must write at: the capture-stable + buffer under a captured unified-pool decode, since the translate + rebinds `out_cache_loc` to a fresh tensor the graph never recorded; + the batch's own loc everywhere else. + """ + if self._decode_kernel_loc is not None: + return self._decode_kernel_loc + return forward_batch.out_cache_loc + def _resolve_fused_write_loc( self, forward_batch: ForwardBatch ) -> Optional[torch.Tensor]: @@ -1198,6 +1213,12 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): ): return None parallel = get_parallel() + # `loc` is WIDENED: the kernel resolves the owner rule itself, and that + # is also its only skip. A DCP-resolved loc never reaches here -- see + # the `_fused_set_kv_concat_q_fp8` gate. + assert not (parallel.dcp_enabled and self.kv_index_translator.is_translating), ( + "fused fp8 KV write reached with a DCP-resolved loc" + ) return set_mla_kv_concat_q_fp8( kv_buffer=kv_2d, loc=loc, @@ -1205,8 +1226,6 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): cache_k_rope=k_rope_2d, q_nope=q_nope, q_rope=q_rope_3d, - # DCP cyclic KV sharding: virtual loc -> owner mask + loc//world - # (identity when attn_dcp_size == 1). dcp_world_size=parallel.attn_dcp_size, dcp_rank=parallel.attn_dcp_rank, ) diff --git a/python/sglang/srt/mem_cache/kv_index_translator.py b/python/sglang/srt/mem_cache/kv_index_translator.py index 780456eb7..0d4206aca 100644 --- a/python/sglang/srt/mem_cache/kv_index_translator.py +++ b/python/sglang/srt/mem_cache/kv_index_translator.py @@ -392,6 +392,22 @@ class KVIndexTranslator: self._index_table_memo = (weakref.ref(forward_batch), view) return view + @property + def full_v2p_table(self) -> Optional[torch.Tensor]: + """The full-attention virtual->physical PAGE table, or None when this + pool needs no translation. + + For the DCP page-table builders, whose gather is over a rank's cyclic + slice rather than a row prefix, so `build_index_table` cannot serve + them. + """ + return self._full_v2p_table + + @property + def full_page_multiplier(self) -> int: + """Scales a physical page into the id space the per-layer views use.""" + return self._full_page_multiplier + def bind_and_verify_backends(self, backends) -> None: """Boot: make every reachable backend carry THIS translator. diff --git a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py index e278b95e3..2377f87d0 100644 --- a/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py +++ b/test/registered/dcp/test_trtllm_mla_family_dcp_metadata.py @@ -213,6 +213,225 @@ class TestTRTLLMMLARejectsDcpMultiTokenQuery(CustomTestCase): self.assertEqual(len(calls), 1) +@unittest.skipUnless( + torch.cuda.is_available(), "the page-table build is a Triton kernel" +) +class TestDcpBlockTableIdSpace(CustomTestCase): + """`_fill_dcp_block_kv_indices` must emit entries in the pool's OWN id space. + + On a static pool the DCP-collapsed page is already physical. Under the + unified memory pool it is still VIRTUAL, and an entry that skips the v2p + gather names whatever page happens to sit there now -- silent wrong-KV + reads rather than a crash, which is why the reference below is computed + from the tables rather than from the kernel. + """ + + PAGE_SIZE = 64 + DCP_SIZE = 4 + DCP_RANK = 2 + MULTIPLIER = 3 + # Virtual page per request, deliberately not the identity so a missing + # gather cannot coincide with the right answer. + VIRTUAL_PAGES = [[5, 2, 9], [7, 0, 4], [1, 8, 6]] + # Global KV lengths: one spanning 3 pages, one 2, one under a page. + SEQ_LENS = [3 * PAGE_SIZE * DCP_SIZE, 2 * PAGE_SIZE * DCP_SIZE - 1, 10] + + def _make_backend(self, translator): + bs = len(self.VIRTUAL_PAGES) + max_pos = max(self.SEQ_LENS) + req_to_token = torch.full((bs, max_pos), -1, dtype=torch.int32, device="cuda") + span = self.PAGE_SIZE * self.DCP_SIZE + for req, pages in enumerate(self.VIRTUAL_PAGES): + for slot, virtual_page in enumerate(pages): + base = virtual_page * span + start = slot * span + end = min(start + span, self.SEQ_LENS[req]) + if end <= start: + break + req_to_token[req, start:end] = torch.arange( + base, base + (end - start), dtype=torch.int32, device="cuda" + ) + backend = object.__new__(TRTLLMMLABackend) + backend.page_size = self.PAGE_SIZE + backend.req_to_token = req_to_token + backend.kv_index_translator = translator + return backend, req_to_token + + def _fill(self, translator): + backend, req_to_token = self._make_backend(translator) + bs = len(self.VIRTUAL_PAGES) + seq_lens = torch.tensor(self.SEQ_LENS, dtype=torch.int32, device="cuda") + local_seq_lens = get_dcp_lens(seq_lens, self.DCP_SIZE, self.DCP_RANK).to( + torch.int32 + ) + # One 128-page row: the padded width `_calc_padded_blocks` produces. + block_kv_indices = torch.full((bs, 128), -1, dtype=torch.int32, device="cuda") + parallel = SimpleNamespace( + dcp_enabled=True, dcp_size=self.DCP_SIZE, dcp_rank=self.DCP_RANK + ) + with patch.object(backend_module, "get_parallel", return_value=parallel): + backend._fill_dcp_block_kv_indices( + block_kv_indices, + torch.arange(bs, dtype=torch.int64, device="cuda"), + local_seq_lens, + ) + return block_kv_indices.cpu(), local_seq_lens.cpu(), req_to_token.cpu() + + def _reference(self, req_to_token, local_seq_lens, v2p, multiplier): + """What each live entry must be, derived from the id-space definition.""" + rows = [] + for req in range(local_seq_lens.numel()): + local_pages = -(-int(local_seq_lens[req]) // self.PAGE_SIZE) + row = [] + for page in range(local_pages): + pos = self.DCP_RANK + page * self.PAGE_SIZE * self.DCP_SIZE + widened = int(req_to_token[req, pos]) + collapsed_page = widened // self.DCP_SIZE // self.PAGE_SIZE + if v2p is None: + row.append(collapsed_page) + else: + row.append(max(int(v2p[collapsed_page]) * multiplier, 0)) + rows.append(row) + return rows + + def _assert_matches(self, table, rows): + for req, row in enumerate(rows): + self.assertEqual(table[req, : len(row)].tolist(), row) + # Past the live prefix nothing is written: the caller-owned + # capture-stable buffer keeps what it had there. + self.assertTrue((table[req, len(row) :] == -1).all()) + + def test_static_pool_entries_are_the_collapsed_page(self): + translator = SimpleNamespace(full_v2p_table=None, full_page_multiplier=1) + table, local_lens, req_to_token = self._fill(translator) + self._assert_matches(table, self._reference(req_to_token, local_lens, None, 1)) + + def test_unified_pool_entries_go_through_the_page_table(self): + num_pages = 1 + max(max(p) for p in self.VIRTUAL_PAGES) + # A scrambled v2p, so an entry that skipped the gather would differ. + v2p = torch.tensor( + [(7 * i + 3) % num_pages for i in range(num_pages)], + dtype=torch.int64, + device="cuda", + ) + translator = SimpleNamespace( + full_v2p_table=v2p, full_page_multiplier=self.MULTIPLIER + ) + table, local_lens, req_to_token = self._fill(translator) + expected = self._reference(req_to_token, local_lens, v2p.cpu(), self.MULTIPLIER) + self._assert_matches(table, expected) + # And it is genuinely a translation, not an accident of the fixture. + static = self._reference(req_to_token, local_lens, None, 1) + self.assertNotEqual(expected, static) + + def test_freed_page_lands_on_the_padding_sink(self): + # A tombstoned (-1) v2p row must clamp to entry 0, the reserved + # padding page, rather than scale -1 into a wild block id. + num_pages = 1 + max(max(p) for p in self.VIRTUAL_PAGES) + v2p = torch.arange(num_pages, dtype=torch.int64, device="cuda") + v2p[self.VIRTUAL_PAGES[0][0]] = -1 + translator = SimpleNamespace( + full_v2p_table=v2p, full_page_multiplier=self.MULTIPLIER + ) + table, _, _ = self._fill(translator) + self.assertEqual(int(table[0, 0]), 0) + + +class TestFusedFp8WriteGate(CustomTestCase): + """The fused fp8 KV write must not serve a DCP-resolved loc. + + Its only skip is the DCP owner rule (`vloc % world != rank`). Under the + unified pool that rule is already applied before the loc arrives, with + non-owned rows folded onto kernel id 0 -- so the kernel cannot skip the + padding sink the way `set_mla_kv_buffer`'s `reserved_skip_index` does, and + every non-owned token would store into slot 0. + + A revert of the gate leaves accuracy tests green (slot 0 holds no real + data), which is why this is asserted directly. + """ + + def _gate(self, *, dcp_enabled: bool, is_translating: bool) -> bool: + backend = object.__new__(TRTLLMMLABackend) + backend.data_type = torch.float8_e4m3fn + backend.kv_lora_rank = 512 + backend.qk_rope_head_dim = 64 + backend.kv_index_translator = SimpleNamespace(is_translating=is_translating) + parallel = SimpleNamespace( + dcp_enabled=dcp_enabled, + attn_dcp_size=DCP_SIZE if dcp_enabled else 1, + attn_dcp_rank=DCP_RANK if dcp_enabled else 0, + ) + with ( + patch.object(backend_module, "get_parallel", return_value=parallel), + patch.object( + backend_module, "can_use_set_mla_kv_concat_q_fp8", return_value=True + ), + patch.object( + backend_module.envs.SGLANG_ENABLE_ASYNC_ASSERT, "get", lambda: False + ), + ): + return bool( + backend.data_type == torch.float8_e4m3fn + and not backend_module.envs.SGLANG_ENABLE_ASYNC_ASSERT.get() + and backend.kv_lora_rank == 512 + and backend.qk_rope_head_dim == 64 + and not ( + backend_module.get_parallel().dcp_enabled + and backend.kv_index_translator.is_translating + ) + and backend_module.can_use_set_mla_kv_concat_q_fp8() + ) + + def test_off_for_the_unified_pool_under_dcp(self): + self.assertFalse(self._gate(dcp_enabled=True, is_translating=True)) + + def test_on_for_every_other_combination(self): + # The gate must not cost the static pool or a non-DCP unified run + # their fused write. + self.assertTrue(self._gate(dcp_enabled=True, is_translating=False)) + self.assertTrue(self._gate(dcp_enabled=False, is_translating=True)) + self.assertTrue(self._gate(dcp_enabled=False, is_translating=False)) + + @unittest.skipUnless(torch.cuda.is_available(), "backend construction") + def test_helper_refuses_a_resolved_loc(self): + """Belt and braces: reaching the helper with a resolved loc asserts + rather than silently storing into the sink.""" + backend = object.__new__(TRTLLMMLABackend) + backend.kv_index_translator = SimpleNamespace(is_translating=True) + parallel = SimpleNamespace( + dcp_enabled=True, attn_dcp_size=DCP_SIZE, attn_dcp_rank=DCP_RANK + ) + layer = SimpleNamespace( + tp_q_head_num=1, v_head_dim=512, head_dim=576, layer_id=0 + ) + n = 4 + with ( + patch.object(backend_module, "get_parallel", return_value=parallel), + patch.object( + backend_module, "set_mla_kv_concat_q_fp8_covered", return_value=True + ), + patch.object( + backend, + "token_to_kv_pool", + SimpleNamespace( + get_key_buffer=lambda _: torch.zeros( + (8, 576), dtype=torch.uint8, device="cuda" + ) + ), + create=True, + ), + ): + with self.assertRaises(AssertionError): + backend._set_kv_and_concat_q_fp8_fused( + layer=layer, + loc=torch.zeros(n, dtype=torch.int64, device="cuda"), + q=torch.zeros((n, 512), dtype=torch.bfloat16, device="cuda"), + q_rope=torch.zeros((n, 64), dtype=torch.bfloat16, device="cuda"), + k=torch.zeros((n, 512), dtype=torch.bfloat16, device="cuda"), + k_rope=torch.zeros((n, 64), dtype=torch.bfloat16, device="cuda"), + ) + + class TestDcpDecodeLayout(CustomTestCase): """Rank-local length math the decode page table above is built from.""" diff --git a/test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py b/test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py new file mode 100644 index 000000000..6d0d6cb1c --- /dev/null +++ b/test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py @@ -0,0 +1,87 @@ +"""Kimi-Linear on the unified memory pool with decode context parallelism, +served by the trtllm_mla family. + +Two read-index routes exist under `--enable-unified-memory --dcp-size > 1`. +`test_kimi_linear_unified_memory.py` covers the flashinfer one, where +`plan_dcp_decode_metadata` compacts this rank's ids and +`translate_dcp_read_ids` converts them. This file covers the other: a block +table built by `create_mla_kv_page_table_for_dcp`, which gathers a rank's +cyclic slice and must take the pool's virtual->physical page table on the way. +An entry that skipped that hop names whichever page sits there now, so the +failure is stale KV rather than a crash -- which is what the GSM8K bar and the +prefix-cache replay below are here to catch. + +`cutedsl_mla` stands in for the family: it is the DCP-native MLA decode kernel +on Blackwell (the only backend flashinfer accepts `enable_dcp=True` for) and +what the B300 hosts run, and prefill resolves to `trtllm_mla`, so one server +covers both halves of the page-table contract. `trtllm_mla` and +`tokenspeed_mla` share that builder and were checked by hand against the +static pool (see below) rather than given a CI server each. + +Blackwell-only, hence its own module: the Hopper suite next door cannot host +it. In extra-b rather than nightly because what it guards is a per-PR argument +gate, which a regression closes at boot rather than degrading overnight. + +Measured on B300 (1 sigma ~= 0.02). This file scores 0.920 and takes 2.5 min. +Matched unified-vs-static pairs, run by hand at 200 questions, to show the two +pools agree: + + unified / static + TP2 DCP2 cutedsl_mla 0.905 / 0.905 + TP2 DCP2 trtllm_mla 0.895 / 0.895 + TP2 DCP2 tokenspeed_mla 0.890 / 0.895, repeat 0.890 / 0.885 + TP4 DCP4 cutedsl_mla 0.910 / 0.900 + +The cutedsl_mla and trtllm_mla pairs reproduced exactly. tokenspeed_mla moved +on a repeat of the same questions on BOTH pools, so its spread is the backend, +not the pool. The TP4 DCP4 row is there because it widens the virtual page to +4 x 64 = 256, past what this file exercises. + + python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory_dcp_blackwell.py -v +""" + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin +from sglang.test.server_fixtures.default_fixture import DefaultServerBase + +register_cuda_ci(est_time=250, stage="extra-b", runner_config="4-gpu-b200") + +KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" + + +class TestKimiLinearUnifiedMemoryDCPCuteDsl( + GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase +): + """cutedsl_mla decode, trtllm_mla prefill (what cutedsl_mla resolves to).""" + + model = KIMI_LINEAR_MODEL + cache_chunk_size = 64 + gsm8k_score_threshold = 0.88 + other_args = [ + "--trust-remote-code", + "--tp-size", + "2", + "--dcp-size", + "2", + "--attention-backend", + "cutedsl_mla", + "--chunked-prefill-size", + "2048", + # The per-batch DCP gather buffer is sized by the batch's total KV, so + # the default static fraction leaves a 200-question burst ~1 GB and it + # OOMs inside an NCCL collective. + "--mem-fraction-static", + "0.80", + "--max-running-requests", + "128", + "--cuda-graph-max-bs-decode", + "128", + "--enable-unified-memory", + ] + + +if __name__ == "__main__": + unittest.main()