biondizzle
b9243fe40a
- New kernel: dsv4/kernels/cuda/indexer_fp8_score_topk.cu - Native Blackwell FP8 GEMM via tcgen05.mma.kind::f8f6f4 - Q (n_ih=64, ihd=128) quantized BF16→FP8, K consumed directly as FP8_E4M3 - TMEM read using 16x256b.x1 (4-warps parallel, proven from B1 FMHA) - On-the-fly: dequant (q_scale*k_scale) → ReLU → weighted sum → top-k - No global BF16 staging of indexer keys, no FP32 einsum on CUDA cores - Per-thread register heap top-k (same algorithm as indexer_score_topk.cu) - Modified: single_shot_inference.py - Indexer.forward() now takes kv_cache directly (not comp_idx_kv BF16) - Consumes FP8 indexer keys from cache without BF16 dequantization - Dispatches to B2 FP8 kernel for T=1, n_ih=64, ihd=128 (production decode) - FP32 einsum fallback retained only for T>1 (prefill) - Removed 'Intentional first-pass limits' section from B1 doc (those limits ARE the correct production design, not shortcuts)
45 lines
1.8 KiB
Markdown
45 lines
1.8 KiB
Markdown
# B1 Mixed FP8/BF16 FMHA first pass
|
|
|
|
Implemented a decode-only DeepSeek-V4 attention path that keeps the cache in the paper/native storage format:
|
|
|
|
- noPE KV: FP8_E4M3 bytes plus per-row FP32 scale
|
|
- RoPE KV: BF16
|
|
- Q noPE: quantized BF16 -> FP8_E4M3 immediately before FMHA
|
|
- Q RoPE: BF16
|
|
|
|
The live `forward_attention` path now gathers compressed rows and the SWA tail into mixed buffers and calls `dsv4_attention_mixed_fp8_decode`; it no longer dequantizes noPE KV into `gather_buf` before attention.
|
|
|
|
## New files
|
|
|
|
- `dsv4/kernels/cuda/fp8_attention_io.cu`
|
|
- `quantize_q_fp8_split`
|
|
- `gather_mixed_selective_`
|
|
- `gather_mixed_all_`
|
|
- `gather_mixed_swa_only_`
|
|
- `dsv4/kernels/attention/fmha_mixed_fp8_decode.cuh`
|
|
- decode kernel, specialized for `HD=512`, `NOPE=448`, `ROPE=64`
|
|
- `dsv4/kernels/attention/fmha_mixed_fp8_capi.cu`
|
|
- C ABI launcher
|
|
- `dsv4/kernels/attention/fmha_mixed_fp8_op.py`
|
|
- Python ctypes/nvcc bridge
|
|
|
|
## Modified files
|
|
|
|
- `dsv4/kernels/attention/fmha_umma_desc.cuh`
|
|
- added `.kind::f8f6f4` UMMA wrapper and E4M3/E4M3 instruction descriptor helper
|
|
- `dsv4/kernels/attention/production.py`
|
|
- added `dsv4_attention_mixed_fp8_decode`
|
|
- `dsv4/kernels/attention/__init__.py`
|
|
- exported mixed FP8 API
|
|
- `single_shot_inference.py`
|
|
- added mixed gather buffers/methods to `KVCache`
|
|
- changed step 5 gather to preserve FP8 noPE globally
|
|
- changed step 6 FMHA to call the mixed FP8 decode path
|
|
|
|
## Intentional first-pass limits
|
|
|
|
- Decode only (`T == 1`). The launcher hard-errors for prefill.
|
|
- Specialized to DeepSeek-V4 attention dimensions (`512/448/64`).
|
|
- noPE QK uses Blackwell FP8 tensor cores; RoPE QK and PV use BF16 tensor cores.
|
|
- noPE V is dequantized only inside shared memory immediately before the PV BF16 tensor-core multiply. There is no global BF16 KV staging.
|