Support Mamba async scheduling

[lmcafee-nvidia]

Summary

• Enable hybrid and Mamba models in the async scheduling path while retaining the existing single recurrent-state bank.
• Preserve each surviving request's Mamba state-slot mapping when async resolution compacts request rows.
• Write greedy sample IDs directly into the reusable integer output buffer so sampling works with both BF16 and FP32 logits.

Architecture

The async-prefill parent drains a pending forward before it admits newly ready work. This gives the existing single Mamba state bank a clear ownership boundary without adding a second bank, admission-specific state, or Mamba kernel changes.

Async request resolution already moves logical request rows after requests finish. This PR extends that move to the Mamba state-slot mapping and clears mappings for vacated rows. The recurrent tensors themselves remain in their existing physical slots.

Performance

DFW inference-bench used the trained hybrid Mamba 2B checkpoint (2.875B parameters), one GSM8K request at a time, OSL 512, and full-iteration CUDA graphs. Each mode ran four requests; the first was discarded as warmup.

Revision / mode	Measured throughput (tok/s)	Steady mean (tok/s)
This PR, legacy	201.4, 201.7, 201.6	201.6
This PR, async	228.1, 229.5, 230.3	229.3
Prior single-bank async implementation	225.4, 225.6, 227.0	226.0

Async scheduling improves steady throughput by 13.8% over legacy and is 1.5% faster than the same-hardware rerun of the prior single-bank implementation.

Validation

• DFW focused inference unit suite: 95 passed.
• DFW trained hybrid Mamba 2B functional test: exact golden match across three staggered requests, including three request-compaction steps.
• DFW trained hybrid Mamba 2B HumanEval: main, legacy, and async all scored 57/164 (0.3476), with byte-identical raw generations.
• DFW trained Nano-v3 Mamba+MoE EP8 HumanEval: main and legacy scored 127/164 (0.7744); async scored 126/164 (0.7683). The one-result launch-to-launch EP variation matches prior qualification of this checkpoint.
• HSG 300M Mamba+MTP+EP4 smoke runs: legacy and async completed four requests with 32 generated tokens. Repeated launches confirmed that observed token variation was launch-level EP numerics rather than scheduling-specific state.

[santhnm2]

Why is this an argmax now instead of a max?

[lmcafee-nvidia]

torch.max was introduced to avoid the original argmax(logits.float()) conversion while writing into preallocated output buffers. however, its values output requires a buffer matching the logits dtype. hybrid/Mamba models can produce either BF16 or FP32 logits, so that buffer is not reliably typed. native-dtype torch.argmax(..., out=sampled_tokens_gpu) avoids the conversion, writes only the required int64 token IDs, and works for both logits dtypes. the parameterized test covers both BF16 and FP32.