[BUGFIX][Mamba][Qwen3.5] Zero freed SSM cache blocks on GPU

[vadiklyutiy]

Essential problem

Fixes #35138
Workaround for Dao-AILab/flash-attention#1974

Hybrid models (e.g. Qwen3.5-397B-A17B) share a unified block pool between attention (fp8/fp16) and Mamba/SSM (fp32) layers. When a block previously used by Mamba (fp32 state) is reallocated to an attention layer with a smaller dtype, leftover fp32 bit patterns can appear as NaN/Inf in the new dtype. Attention kernels (FlashAttn3, FlashInfer-TRTLLM, etc.) use multiply-by-zero masking for unused positions, which does not clear NaN (0 * NaN = NaN). The stale NaN then propagates across all requests sharing the same KV-cache block, causing progressive accuracy degradation over time.

What this PR does

Zeroes GPU memory of freshly allocated full-attention KV-cache blocks before they are used, but only for hybrid models (models with Mamba layers). Mamba/SSM blocks are not zeroed (they overwrite their state fully on each step). The approach:

1. Scheduler side — SingleTypeKVCacheManager tracks block IDs allocated since the last scheduling step (only for FullAttentionSpec layers). After scheduling, the scheduler drains these IDs into SchedulerOutput.new_block_ids_to_zero, gated behind self.has_mamba_layers.

2. Worker side — GPUModelRunner._update_states() receives the block IDs and calls _zero_block_ids(), which launches a single Triton kernel (_zero_kv_blocks_kernel) to zero the corresponding memory across all KV-cache segments in one GPU launch.

3. Optimized zeroing — A one-time _init_kv_zero_meta() precomputes absolute byte addresses of all KV-cache segments (handling both block-dim-0 and block-dim-1 layouts, multi-buffer backends, and virtual block splitting). Block IDs are transferred via pre-allocated pinned memory to overlap the H2D copy with kernel launch. This avoids 15 separate index_fill_ calls (For Qwen3.5-379B, one per layer).

4. CuMem compatibility — _init_kv_zero_meta() is called in gpu_worker.py outside the CuMem pool context, so the bookkeeping tensors (segment addresses, block-ID buffers) use the standard PyTorch allocator and survive sleep/wake cycles.

Scope

• Hybrid models only — The zeroing is gated by self.has_mamba_layers in the scheduler. Non-hybrid (pure attention) models are completely unaffected.
• Only freshly allocated blocks — Prefix-cached blocks (cache hits) are not zeroed, preserving prefix caching correctness and performance.
• Only FullAttentionSpec blocks — Mamba/SSM blocks are not zeroed (they overwrite their state fully each step); only attention blocks that may inherit stale Mamba fp32 data are cleared.

Performance overhead

Measured on B200, Qwen/Qwen3-0.6B, BS=500:

Phase	Blocks zeroed	Median latency	vs. forward step
Prefill (BS ~8K)	~515 blocks (~920 MiB)	~170 μs	~1% of 18ms step
Decode (BS ~500)	~30 blocks (~55 MiB)	~15 μs	~0.1% of 13ms step

End-to-end benchmark on B200 (Qwen3.5-397B-A17B-FP8, TP=1 PP=1 DP=8, 2048 prompts, 500 output tokens) showed no measurable throughput degradation (output tokens/s within ±2% noise).

Test plan

• The original issue ([Bug]: Qwen/Qwen3.5-397B-A17B-FP8 and Qwen/Qwen3.5-397B-A17B has accuracy issues when running with Flashinfer Attention backend on Blackwell. #35138) reproduction test passes: 5 consecutive GSM8K evaluations on Qwen3.5-397B-A17B-FP8 (TP=4, FlashInfer backend) — accuracy stable at ~90% across all runs with no degradation. Previously, accuracy would drop significantly after the first run due to NaN accumulation.
• End-to-end throughput benchmark (B200, 2048 prompts) confirms no performance regression.

[gemini-code-assist]

Code Review

This pull request introduces a bug fix for Mamba-based models, specifically addressing an issue where freed SSM cache blocks on the GPU were not being zeroed out. This could lead to incorrect state being reused in subsequent computations. The fix involves implementing a mechanism to track SSM blocks that are truly freed (i.e., their reference count drops to zero) during each scheduling step. These freed block IDs are then passed to the worker, which explicitly zeroes out the corresponding state tensors on the GPU. The changes are well-contained and correctly implemented across the scheduler and worker components, ensuring that Mamba's stateful cache is properly managed. The logic for identifying and collecting freed blocks is soundly integrated into the existing KV cache management lifecycle methods.

[LucasWilkinson]

Could you please provide data on possible perf overhead? also with async scheduling I think it may be risky to zero on free, we may need to move this into the model runner to ensure it ends up in the correct order in the GPU stream

[vadiklyutiy]

we may need to move this into the model runner to ensure it ends up in the correct order in the GPU stream

Actual zeroing happens in GPUModelRunner._update_states. "On free" we just collect corresponding blocks.
Or didn't get you comment?

[vadiklyutiy]

Below is from slack discussion. But I think it's worth attention and sharing here

Seems both FlashAttn and trt-llm attn use mul by 0 to mask not used values.
To be sure that its correct we must guarantee that no NaN. In general case masking by mul by 0 is incorrect.
In my flavor, it should/must be fixed in kernels.
Seems even common full attn can produce NaN for some corner cases. It is likely ok to produce NaN(and garbage tokens) for one specific request, but with kernels that not tolerant to NaN, we propagate this problem to another good requests with subsequent garbage for all requests.

[tdoublep]

How does this interact with prefix caching? If we zero out blocks when their ref_cnt hits zero, doesn't that mean they can't be re-used if something comes along later that gets a cache hit?

Wouldn't it to be better to detect the event when we use a block for attention that was previously used for mamba (in some other dtype) and zero it out at that point?

[vadiklyutiy]

Could you please provide data on possible perf overhead? also with async scheduling I think it may be risky to zero on free, we may need to move this into the model runner to ensure it ends up in the correct order in the GPU stream

Ran on B200

VLLM_USE_FLASHINFER_MOE_FP8=1 vllm serve Qwen/Qwen3.5-397B-A17B-FP8 --port 8000 -tp 1 -pp 1 -dp 8 --enable-expert-parallel --language-model-only --reasoning-parser qwen3 --kv-cache-dtype fp8 --stream-interval=100

 vllm bench serve --backend vllm --model Qwen/Qwen3.5-397B-A17B-FP8 \
            --endpoint /v1/completions --dataset-name random --random-input 2 \
            --random-output 500  --max-concurrency 512 --num-prompt 2048 \
            --ignore-eos --temperature=0.0 

With changes and without changes the Output Tokens vary around 15000+-300. Definitely nothing dramatical from perf point of view but not exact numbers.

[vadiklyutiy]

How does this interact with prefix caching? If we zero out blocks when their ref_cnt hits zero, doesn't that mean they can't be re-used if something comes along later that gets a cache hit?

Wouldn't it to be better to detect the event when we use a block for attention that was previously used for mamba (in some other dtype) and zero it out at that point?

Good catch I broke the prefix caching :/
Move to draft to think better way

[vadiklyutiy]

How does this interact with prefix caching? If we zero out blocks when their ref_cnt hits zero, doesn't that mean they can't be re-used if something comes along later that gets a cache hit?

Wouldn't it to be better to detect the event when we use a block for attention that was previously used for mamba (in some other dtype) and zero it out at that point?

I redo this PR similar to what @tdoublep proposed.

I decided to zero out every new block, whether it comes from attention or from the SSM.

Justification: attention can also produce NaNs in certain corner cases. Getting garbage for one specific request is likely acceptable, but without zeroing, the NaN could propagate to all requests.

[vadiklyutiy]

pls take a look

[benchislett]

Would it be more efficient to build an index tensor and have one op to zero at all the block id slots?

[benchislett]

How many block ids would we normally see for a typical prefill/decode? Is it very few?

[vadiklyutiy]

This zeroing takes small amount of time. We do it once per forward step and only for new.

@benchislett Can you say right away does it code works in sync or async part?

[benchislett]

It happens always.

[benchislett]

I notice that this is not specific to SSM blocks, and it clears all new KV blocks. Will this have a detrimental effect on prefills for non-mamba deployments where block_size=16?

In this case if we get a prefill of 8k tokens, that will be 512 new blocks, right? I think that would lead to 512 kernel invocations in this implementation. If that is indeed the case, this will not suffice.

[vadiklyutiy]

right,
I am optimizing it

[pavanimajety]

does it make sense to use torch.tensor for block ids and use a gpu operation to zero the indices in tensors?

[vadiklyutiy]

See my comment below

[vadiklyutiy]

does it make sense to use torch.tensor for block ids and use a gpu operation to zero the indices in tensors?

I implemented zeroing as a triton kernel

[vadiklyutiy]

Pls lets me know if there is a better way to do it