amd/deepseek_v4 integration 4/N - TilelangAttn 0428#24033
Conversation
1am9trash
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
hebiao064,
hlu1,
hubertlu-tw,
ispobock,
kkHuang-amd and
merrymercy
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request replaces the BF16 sparse attention kernel with a more efficient 2-stage FP8 attention kernel system consisting of partial and combine kernels. The new implementation supports in-kernel dequantization and dual KV caches. Additionally, it includes a workaround for a tilelang mutation bug and utility functions for cache reinterpretation. Review feedback focuses on improving code maintainability by replacing magic numbers with constants, using more idiomatic PyTorch APIs for integer limits, and addressing the use of deprecated storage methods.
| def _build_fp8_combined_view(k_cache: torch.Tensor) -> Tuple[torch.Tensor, int, int]: | ||
| """ | ||
| Reinterpret a MODEL1_FP8Sparse KV cache as a contiguous uint32 view. | ||
| Input: k_cache (num_blocks, block_size, 1, d_qk) fp8/uint8 | ||
| — per-block storage also holds scales + padding past d_qk. | ||
| Output: (num_blocks, block_pad_u32) uint32 covering the full block | ||
| stride. Same storage ashe input, no copy. | ||
| """ | ||
| k_u8 = k_cache.view(torch.uint8) if k_cache.dtype != torch.uint8 else k_cache | ||
| num_blocks = k_u8.shape[0] | ||
| block_size = k_u8.shape[1] | ||
| block_pad_u32 = k_u8.stride(0) // 4 | ||
| storage = k_u8.untyped_storage() | ||
| flat_u32 = torch.empty(0, dtype=torch.uint32, device=k_u8.device).set_( | ||
| storage, 0, (storage.nbytes() // 4,), (1,) | ||
| ) | ||
| k_combined = torch.as_strided( | ||
| flat_u32, | ||
| size=(num_blocks, block_pad_u32), | ||
| stride=(block_pad_u32, 1), | ||
| storage_offset=k_u8.storage_offset() // 4, | ||
| ) | ||
| return k_combined, num_blocks, block_size |
There was a problem hiding this comment.
The function _build_fp8_combined_view uses untyped_storage() which is deprecated in newer PyTorch versions. Consider using storage() or untyped_storage() if you are sure about the version compatibility, but untyped_storage() is generally discouraged in favor of storage() for newer codebases.
| return max(tilelang.math.next_power_of_2(head_kv), 16) | ||
| return max(_next_power_of_2(head_kv), 16) | ||
| _TOPK_LEN_SENTINEL_CACHE: dict = {} | ||
| _INT32_MAX = 2**31 - 1 |
There was a problem hiding this comment.
Using 2**31 - 1 is fine, but torch.iinfo(torch.int32).max is more idiomatic and readable for defining the maximum value of a 32-bit integer.
| cur = _TOPK_LEN_SENTINEL_CACHE.get(device) | ||
| if cur is None or cur.numel() < batch: | ||
| cur = torch.full( | ||
| (max(batch, 256),), _INT32_MAX, dtype=torch.int32, device=device |
There was a problem hiding this comment.
The use of max(batch, 256) for the sentinel tensor size is a magic number. It should be defined as a named constant at the module level to improve maintainability.
2 participants
Motivation
Update amd/deepseek_v4 integration branch
Following PRs have large set of conflict, we use this PR and upstream amd/deepseek_v4 branch to integrate in parallel.
#23600
#23608
The original
flash_mla_with_kvcache_torchlaunches 101 kernels per call; switching to the tilelang kernel reduces this to 2.Controlled by
export SGLANG_HACK_FLASHMLA_BACKEND=tilelang.Modifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ci