fix(mamba): separate state and checkpoint chunk sizes

[momaekar1]

Important

This PR contains only the Mamba checkpoint semantic fix. It does not include
#32500 or #35013. Merge this after #32500; #34560 is already merged into
main. The end-to-end validation stack also included #35013 for packed MTP
HiCache transfers.

Motivation

Mamba radix caching currently uses mamba_cache_chunk_size for two different
concepts:

• the page-aligned granularity at which checkpoints can be stored in the radix
  cache; and
• the physical chunk size used by the kernel to pack the intermediate state
  tensor h.

These values are not always equal on Ascend. Qwen3.6 GDN uses 64-token kernel
state chunks with page_size=128, so L1 checkpoints remain on a 128-token grid
while h is packed every 64 tokens. Indexing h with the L1 grid restores the
wrong recurrent state. The incorrect state is selected before tier transfer,
so L2 and L3 can both replay the same wrong checkpoint; this is not a Mooncake
transport issue.

CUDA does not use this new semantic. The shared backend and non-NPU scheduler
path continue to use the existing mamba_cache_chunk_size, preserving their
current behavior.

This is the minimal checkpoint-semantics part of #33515, adapted to the latest
main. It intentionally does not include that PR's cold-prefill boundary
change.

Minimal reproduction

Use Qwen3.6 GDN on Ascend with:

kernel state chunk = 64
L1 page/checkpoint grid = 128
extend length = 3008
checkpoint depth = 2944

The scheduler uses a synthetic tracked length of 2945 to select the interior
state at checkpoint 2944. The packed state index must therefore be:

2944 / 64 = h[46]

The previous code instead used the L1 checkpoint grid and selected:

2944 / 128 = h[23]

The NPU backend comment records this concrete L1 failure, and the regression
test covers the per-request packed offset with expected h_src=[46, 93] for a
two-request batch.

Modifications

• Add npu_mamba_state_chunk_size for the Ascend kernel's packed intermediate
  states.
• Override the state chunk only in AscendMambaAttnBackendBase.
• Use the new state chunk in scheduler tracking only when running on NPU.
• Keep mamba_cache_chunk_size and the checkpoint grid unchanged for radix
  placement, branching, convolution tracking, and all non-NPU backends.
• Preserve the latest mamba_checkpoint_grid/DCP checkpoint-depth behavior.
• Add focused regression tests for the 64-token state / 128-token L1 case and
  for unchanged shared-backend behavior.

L1 divergence comparison

The A/B test used commits that differ only by this PR's patch. The before commit
was the validation stack parent c2f86df61; the after commit 14f0e87f3 has
the same stable patch-id as this PR head. Both runs used Qwen3.6-27B BF16, TP2,
Ascend NPU, NEXTN, 4 prompts with 3000 input tokens and 512 output tokens, and
reused exactly 2944 L1 device tokens for every prompt.

                         before PR              after PR
first mismatch           [65, 88, 101, 23]      [217, 322, 293, 58]
weighted accept rate     60.4396%                62.8169%
weighted accept length   2.8132                  2.8845

The cold baseline was identical in both runs: weighted accept rate 63.4988%
and weighted accept length 2.9050. The later first mismatches and improved
NEXTN statistics show that the wrong h index is fixed. Replay is still not
token-identical to cold generation; that remaining divergence is the separate
cold-prefill boundary issue intentionally left out of this PR.

L2/L3 replay verification

The end-to-end validation stack was latest main with #34560,
#32500@9acc6b733, #35013, and this PR. Model and request settings matched the
L1 test above.

• PR-focused tests on the standalone PR head: 6 passed.
• Focused scheduler, state-index, and stacked NPU host-pool tests: 8 passed.
• Host L2: host=2944, device=0, storage=0 for 4/4 prompts.
• Mooncake L3: storage=2944, device=0, host=0 for 4/4 prompts.
• L2 and L3 replay output IDs were identical for all four prompts:
  sha256=a0349ca2a8cdf023da0ef7b86370fbc25eac55ba5d8a656665a845c4a1d969d9.
• L2 and L3 speculative statistics were identical: weighted accept rate
  62.8169%, weighted accept length 2.8845, and correct/proposed drafts
  1338/2130.

L2 and L3 had the same first mismatch positions as the fixed L1 run:
[217, 322, 293, 58]. Mooncake therefore adds no additional divergence.

Accuracy verification

Ran the repository's benchmark/gsm8k/bench_sglang.py against the same full
NEXTN + Mooncake validation stack using the first 200 examples from the
official GSM8K test set:

num_questions=200
num_shots=5
parallel=4
temperature=0
max_new_tokens=512
chat_template_kwargs={"enable_thinking": false}

Result: 194/200 correct (97.0% accuracy), 0% invalid responses, no request
failures, and 218.519 s latency. The previous incomplete-stack result was
195/200; only prompt 182 changed correctness, from the correct answer 23 to
24.

Speed tests and profiling

This change adds no kernel launch or state copy. It only gives the Ascend state
index calculation its own chunk semantic. The 200-example accuracy run took
218.519 s; this is not a controlled performance comparison.

Checklist

• Format code with pre-commit.
• Add focused unit tests, including unchanged non-NPU semantics.
• Compare L1 replay before and after the exact PR patch.
• Complete Qwen3.6 L2/L3 replay validation on the stated stacked setup.
• Complete the 200-example GSM8K accuracy rerun on the stated stacked setup.
• Release NPU resources after validation.
• Documentation is not required for this internal correctness fix.

---

CI States

Latest PR Test (Base): ⏳ Run #32013403511
Latest PR Test (Extra): ⏳ Run #32013403405