optimize the compressor 128 split cutedsl kernel

[Jie-Fang]

Purpose

This PR optimizes the compressor 128 split cutedsl kernel which utilizes the block_size=8 of state_cache for dpsk v4.

Summary

This PR optimizes the DeepSeek V4 C128 CuTeDSL sparse-attention compressor path.

The optimized path is intentionally scoped to the real vLLM C128 layout:

• head_size = 512
• state_width = 512
• compress_ratio = 128
• overlap = False
• compressor state-cache block size = 8
• paged KV-cache block size is treated as a compile-time static value

For C4, this PR keeps using the existing fused CuTeDSL kernel path, since that is the path used in real vLLM inference and is faster than the split version for C4.

Changes

• Replace the generic C128 split compressor with SparseAttnCompressC128Block8Kernel, specialized for the real DeepSeek V4 C128 state-cache layout.
• Keep C4 on the fused SparseAttnCompressNormRopeStoreC4Kernel path.
• Make the C128 store kernel use a compile-time static KV-cache block size.
• Add wrapper-side validation that kv_cache_block_size == k_cache.shape[1].
• Reject unsupported C128 layouts early instead of silently running with assumptions that would produce incorrect results.

Why This Is Faster

The new C128 split compressor takes advantage of the real vLLM C128 layout where one C128 compression window maps to 16 state-cache rows with block size 8.

Compared with the generic split kernel, it avoids per-row dynamic block lookup and dynamic offset calculation. Each warp directly loads two consecutive state blocks from block_table, processes 16 rows, and reduces across 8 warps for each 64-wide head tile. This reduces indexing overhead and improves memory access/reduction efficiency for the C128 decode-boundary path.

Validation

Tested on NVIDIA B300.

Baseline for performance comparison: latest upstream/main at f91fb2fcf.

Passed:

• pre-commit run --files vllm/models/deepseek_v4/nvidia/ops/sparse_attn_compress_cutedsl.py
• python -m pytest tests/kernels/test_compressor_kv_cache.py -q
• tests/kernels/test_compressor_kv_cache.py: 32 passed

Additional correctness coverage:

• C128 decode-boundary, real state block size 8 / KV block size 2: byte-equal vs Triton fused output
• C128 varlen real-layout scenario: byte-equal vs Triton fused output
• C128 mixed boundary/inactive-token scenario: byte-equal vs Triton fused output; inactive KV slots remain unchanged
• C4 fused decode-boundary smoke: byte-equal vs Triton fused output

Representative correctness results:

Scenario	max_dequant_diff vs Triton	max_rope_diff vs Triton	raw_nope_equal	raw_scale_equal	Notes
C128 decode-boundary	0	0	True	True	compress ref diff 3.35e-08
C128 varlen real layout	0	0	True	True	compress ref diff 1.49e-07
C128 mixed boundary/inactive	0	0	True	True	3 inactive KV slots unchanged
C4 fused smoke	0	0	True	True	fused path validation

Performance

Measured on B300 with warmup=50, iters=1000, bench_inputs=4.

token_num	this PR us	upstream/main us	speedup
1	18.869	20.469	1.08x
2	18.762	20.313	1.08x
4	18.780	20.374	1.08x
8	18.813	20.480	1.09x
16	18.757	20.414	1.09x
32	18.735	20.342	1.09x
64	18.592	20.426	1.10x
128	18.564	29.765	1.60x
256	29.048	47.512	1.64x
512	51.199	86.584	1.69x
1024	95.178	160.741	1.69x
2048	183.418	309.586	1.69x
4096	361.789	607.595	1.68x
8192	709.927	1200.419	1.69x
16374	1425.900	2390.316	1.68x

---

Essential Elements of an Effective PR Description Checklist

• [ x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• [ x] The test plan, such as providing test command.
• [x ] The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Jie-Fang.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @Jie-Fang, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10