[Bugfix][Kimi K3] Split DSpark draft out of the target MLA KV group

[lucifer1004]

What

The K3-native DSpark draft marks its MLA layers non_causal_multi_token_decode=True. MLAAttentionSpec.merge ORs that flag, so when the draft's spec merges with the target's identically-shaped MLA spec, the causal target group inherits it too. That raises the group's TritonMLA reorder threshold and misroutes the target's short prefills and causal verification blocks into a decode path that expects one query row per request — producing corrupted output and illegal memory accesses during CUDA graph capture on Kimi-K3 DSpark.

Tag the draft's spec with model_version="kimi_k3_dspark" so the KV-cache group splitter keeps it apart from the target (merge requires a uniform model_version).

Not a duplicate

• Distilled from [DSpark] Support pipeline-parallel targets in disaggregated and aggregated serving (+ padded graph batch safety) #53577 (split requested in review). [Bugfix] Annotate MTP draft KV cache groups positionally on the hybrid grouping path #55390 annotates MTP draft KV groups on the hybrid grouping path — different mechanism (positional annotation vs. model-version split) and not K3-DSpark-specific.

Tests

• New tests/models/kimi_k3/test_dspark_kv_group.py (3 tests): the draft spec carries the model version, the target spec does not, and MLAAttentionSpec.merge refuses to combine them.
• python -m pytest tests/models/kimi_k3/test_dspark_kv_group.py -q → 3 passed.

Model evaluation

Verified end-to-end as part of #53577: Kimi-K3 DSpark native-draft TP8 and PP2×TP4 serving previously produced corrupted output / capture IMA without this split and are clean with it.

AI assistance was used in preparing this PR.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[GirasoleY]

I feel we may still want to allocate the kv cache as one group, but fix trtion mla behavior.

A quick search in community it feels #51065 is closer to a proper fix. Could you check if this PR fix your use case?

[lucifer1004]

@GirasoleY Confirmed — I tested #51065 against this exact use case (Kimi-K3 + K3-native DSpark, TP8, SM120, TritonMLA path, no KV-group split): CUDA graph capture is clean (the IMA is gone), generations are coherent (the corruption is gone), and speculative decoding is actively accepting drafts. Agreed the backend fix is the proper root-cause fix; closing this PR in favor of #51065. Details posted in #51065 (comment)