B1: mixed FP8/BF16 decode FMHA integration
- New: fmha_mixed_fp8_decode.cuh (Blackwell FP8 tensor-core FMHA kernel) - New: fmha_mixed_fp8_capi.cu (C ABI launcher) - New: fmha_mixed_fp8_op.py (Python ctypes/nvcc bridge) - New: fp8_attention_io.cu (Q quantize + mixed KV gather kernels) - New: fmha_umma_desc.cuh additions (f8f6f4 UMMA + idesc helpers) - Modified: production.py (dsv4_attention_mixed_fp8_decode API) - Modified: single_shot_inference.py (B1 gather + FMHA path) - Modified: __init__.py (export mixed FP8 API) - New: docs/B1_MIXED_FP8_FMHA.md, FINAL_STRETCH.md noPE KV stays FP8_E4M3 + per-row scale, RoPE stays BF16. No global FP8->BF16 KV staging before FMHA. Decode-only (T==1), specialized HD=512/NOPE=448/ROPE=64. CUDA compile/runtime validation pending on B200.
This commit is contained in:
55
docs/B1_MIXED_FP8_FMHA.md
Normal file
55
docs/B1_MIXED_FP8_FMHA.md
Normal file
@@ -0,0 +1,55 @@
|
||||
# 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.
|
||||
|
||||
## Validation status
|
||||
|
||||
The sandbox used to make this patch does not have `nvcc`, so CUDA compilation/runtime validation was not possible here. Python syntax was checked with:
|
||||
|
||||
```bash
|
||||
python3 -m py_compile single_shot_inference.py \
|
||||
dsv4/kernels/attention/production.py \
|
||||
dsv4/kernels/attention/fmha_mixed_fp8_op.py
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user