[Revert] Restore per-group K3 Mamba metadata preparation

[khluu]

Summary

Revert #52388 and restore the prior per-group Kimi K3 Mamba aligned-state-index preparation.

Main build #85529 reproduced the same CUDA illegal-address failure twice in test_can_initialize_large_subset[KimiK3ForConditionalGeneration] on two different physical H200 hosts. A later unreverted main, #85541, reproduced it again on a third physical H200 host. All three failures point at the new get_aligned_state_indices_multi_group_kernel during the second CUDA-graph capture. The cached raw block-table pointers are initialized during the temporary CUDA-graph memory-profile capture and remain cached after that temporary allocation is released, so the later real capture can dereference stale addresses.

• #85529 original: https://buildkite.com/vllm/ci/builds/85529#01a039d3-160f-48d3-8146-9b807dfdf6a2 (h200-ci-4)
• #85529 unchanged retry: https://buildkite.com/vllm/ci/builds/85529#01a03a16-dabc-453c-be5f-e038876c6798 (h200-ci-2)
• #85541 current-tip recurrence: https://buildkite.com/vllm/ci/builds/85541#01a03a23-80cf-4c73-8040-7336c53622df (h200-ci-6)

Why revert

This returns to the exact pre-#52388 implementation while a safe fused-pointer lifetime design is worked out. The affected PR's own CI #85384 passed the B200 Kimi job and normal H200 initialization job, but H200 Basic Models (Extra Initialization) Shard 1 was blocked, so the exact failing test did not run before merge. Current-tip PR #53290's exact-head #85358 did select shard 1 and failed the same exact Kimi test; the duplicate shard status name surfaced a passing sibling in GitHub and masked that shard failure.

Duplicate-work check

Searched open vLLM PRs for get_aligned_state_indices_multi_group_kernel, Kimi K3, and CUDA illegal address; no existing fix or revert covers this failure. #53766 changes only metadata unit-test setup for a separate assertion and does not change the failing runtime path.

Validation

• git diff 41729fc53b02cb09beda2a1ced690d021443be43^ -- <four reverted files> --exit-code: passed; this is the exact inverse of [K3 Perf] Optimize k3 mamba metadata preparation, 6.6~7.6x kernel performance improvement #52388
• git diff --check: passed
• uv tool run pre-commit run --files tests/models/kimi_k3/test_kda_metadata.py vllm/models/kimi_k3/nvidia/kda_metadata.py vllm/v1/worker/gpu/model_states/mamba_hybrid.py vllm/v1/worker/mamba_utils.py: passed
• Exact-head Buildkite #85545 at 063794b877e9cb60c1dedeff7227ec228fd790e8 passed all four H200 Basic Models (Extra Initialization) shards: https://buildkite.com/vllm/ci/builds/85545
• Decisive shard 1 passed models/test_initialization.py::test_can_initialize_large_subset[KimiK3ForConditionalGeneration] and the full selected shard: 90 passed / 3 deselected / 106 warnings in 43m53s. Both PIECEWISE graph-capture phases and FULL capture completed without CUDBG_EXCEPTION_WARP_ILLEGAL_ADDRESS: https://buildkite.com/vllm/ci/builds/85545#01a03a29-74d3-415a-a8f5-49b8cbcf3b05

The exact pre-merge revert gate is green. Human review remains required; after merge, incident closure requires the exact shard to pass on post-merge main.

AI assistance was used to investigate the CI failure and prepare this draft. A human submitter must review every changed line and the evidence before marking it ready.

[khluu]

Exact H200 validation is running in Buildkite #85545 at the literal head 063794b877e9cb60c1dedeff7227ec228fd790e8, selecting only basic-models-tests-extra-initialization; shard 1 is the decisive gate. This supersedes #85544, which was canceled before useful execution because its build request used an abbreviated commit string; no test ran there.

[khluu]

Validation update: #85545 bootstrap passed and the exact H200 Basic Models (Extra Initialization) Shard 1 rendered correctly. It is waiting behind the image dependency; no test has run yet.

[khluu]

Exact validation has started: Buildkite #85545 H200 Extra Initialization shard 1 began at 18:40 UTC on literal head 063794b877e9cb60c1dedeff7227ec228fd790e8. The exact test_can_initialize_large_subset[KimiK3ForConditionalGeneration] is selected. The shard is running earlier models; the Kimi case has no result yet. Merge remains blocked on the exact case and full shard passing.

[khluu]

Exact-head validation update for 063794b877e9cb60c1dedeff7227ec228fd790e8: the targeted Kimi K3 initialization test passed in Buildkite #85545, H200 Extra Initialization shard 1. models/test_initialization.py::test_can_initialize_large_subset[KimiK3ForConditionalGeneration] completed both PIECEWISE capture phases and FULL capture without CUDBG_EXCEPTION_WARP_ILLEGAL_ADDRESS. The full 90-test shard is still running, so I am withholding a merge recommendation until the whole shard exits 0.

[khluu]

Exact-head validation is terminal green at 063794b877e9cb60c1dedeff7227ec228fd790e8.

• Buildkite #85545 passed all four selected H200 Basic Models (Extra Initialization) shards.
• Decisive shard 1 passed models/test_initialization.py::test_can_initialize_large_subset[KimiK3ForConditionalGeneration] and the full shard: 90 passed / 3 deselected / 106 warnings in 43m53s. Both PIECEWISE graph-capture phases and FULL capture completed without CUDBG_EXCEPTION_WARP_ILLEGAL_ADDRESS.
• The unreverted current-tip main #85541 shard 1 independently failed the same exact Kimi case on a third physical H200 host with the same illegal-address kernel signature (1 failed / 89 passed / 3 deselected).

This supplies the exact pre-merge revert gate. The PR remains draft/review-required; after human review, I recommend merging the revert. Incident closure still requires the exact shard to pass on post-merge main.

— Sherlock, Kevin Luu’s CI-monitoring agent

[khluu]

Fresh unreverted-main recurrence: full-daily #85561 selected the exact H200 Extra Initialization shard at main 80771bbbddf9.

• Exact job: https://buildkite.com/vllm/ci/builds/85561#01a03aba-6c4b-415f-8da6-8ff899541570
• Exact case: models/test_initialization.py::test_can_initialize_large_subset[KimiK3ForConditionalGeneration]
• Result: first PIECEWISE capture and FULL capture completed, then the second PIECEWISE capture raised CUDBG_EXCEPTION_WARP_ILLEGAL_ADDRESS on warp 2 in get_aligned_state_indices_multi_group_kernel; the case failed at 21:16:03 UTC.
• Host: h200-ci-3, making this the same deterministic signature on a fourth physical H200 host after h200-ci-4, h200-ci-2, and h200-ci-6.

Culprit #52388 exact-head #85384 did not select this shard, so no exact test ran there. The latest-merge PR #51292 exact-head #85498 also did not select this shard; no test ran, and its batch-invariance all-reduce scope is unrelated.

The revert's literal-head #85545 exact shard remains fully green (90 passed / 3 deselected), but this PR is still draft/review-required. Merge plus an exact post-merge main pass remains the closure gate.

— Sherlock, Kevin Luu’s CI-monitoring agent