diff --git a/README.md b/README.md index 52d4ed750c..4b51e65ec1 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,13 @@ Demo: [4× V100 running Qwen3.8-27B-NVFP4-DFlash2](https://www.bilibili.com/vide # 📊 Performance First +SM70 Flash-V100 now resolves `--kv-cache-dtype fp8` to E4M3. DFlash2 E4M3 +verification uses repaired FP32 attention state, and the Qwen3.8 DFlash2 +configuration enables FP32 logits by default. Rebuild Flash-V100 for precision +revision 4; see [the precision contract and validation](docs/design/sm70_dflash2_fp32_defaults.md). +Historical E5M2/FP16-partial performance results below keep their original +configuration and are not speed claims for these precision defaults. + ## Long-Context Attention: 17.92 → 47.1 → ≈60.8 TFLOP/s | Stage | Evidence | Useful causal Attention compute | Notes | diff --git a/benchmarks/benchmark_sm70_dflash2_fp32_attention.py b/benchmarks/benchmark_sm70_dflash2_fp32_attention.py new file mode 100644 index 0000000000..73de629f06 --- /dev/null +++ b/benchmarks/benchmark_sm70_dflash2_fp32_attention.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compare legacy E4M3 partials with repaired FP32 attention, not FP8 loss.""" + +import argparse +import hashlib +import json +from pathlib import Path + +import torch +from flash_attn_v100 import ( + flash_attn_grouped_e4m3_fp32_paged, + flash_attn_grouped_verify_paged, +) +from flash_attn_v100.flash_attn_interface import flash_attn_v100_cuda + + +def measure(length, page, samples): + torch.manual_seed(20260908) + pages = (length + page - 1) // page + capacity = pages * page + q = torch.randn((8, 6, 256), dtype=torch.float16, device="cuda") + raw = torch.randn((2, capacity, 1, 256), dtype=torch.float16, device="cuda") + encoded = raw.to(torch.float8_e4m3fn).view(torch.uint8) + backing = torch.empty((pages, 2, page, 1, 256), dtype=torch.uint8, device="cuda") + k, v = backing.unbind(1) + order = torch.randperm(pages, device="cuda") + k[order] = encoded[0].reshape_as(k) + v[order] = encoded[1].reshape_as(v) + table = order.int()[None].contiguous() + seq = torch.tensor([length], dtype=torch.int32, device="cuda") + rows = torch.arange(length - 7, length + 1, dtype=torch.int32, device="cuda") + outputs = [torch.empty_like(q), torch.empty_like(q)] + + def legacy(): + flash_attn_grouped_verify_paged( + q, + k, + v, + table, + seq, + out=outputs[0], + softmax_scale=0.0625, + kv_cache_dtype="fp8_e4m3", + k_scale=0.5, + v_scale=1.25, + one_pass=True, + ) + + def repaired(): + flash_attn_grouped_e4m3_fp32_paged( + q, + k, + v, + table, + rows, + out=outputs[1], + softmax_scale=0.0625, + k_scale=0.5, + v_scale=1.25, + ) + + graphs = [] + for call in (legacy, repaired): + for _ in range(20): + call() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + graphs.append(graph) + timings = [[], []] + for _ in range(samples): + for index in (0, 1, 1, 0): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + graphs[index].replay() + end.record() + end.synchronize() + timings[index].append(start.elapsed_time(end)) + + # FP64 resolves the FP16 output-rounding floor; also report the requested + # PyTorch FP32 oracle on exactly the same quantized KV and causal rows. + errors = [{}, {}] + for dtype in (torch.float32, torch.float64): + rk = encoded[0, :length, 0].view(torch.float8_e4m3fn).to(dtype) * 0.5 + rv = encoded[1, :length, 0].view(torch.float8_e4m3fn).to(dtype) * 1.25 + scores = q.transpose(0, 1).to(dtype) @ rk.T * 0.0625 + mask = torch.arange(length, device="cuda")[None] >= rows[:, None] + scores.masked_fill_(mask[None], -torch.inf) + expected = (scores.softmax(-1) @ rv).transpose(0, 1) + denominator = expected.norm() + floor = float((expected.half().to(dtype) - expected).norm() / denominator) + for error, output in zip(errors, outputs): + diff = output.to(dtype) - expected + error[str(dtype)] = { + "relative_l2": float(diff.norm() / denominator), + "max_abs": float(diff.abs().max()), + "fp16_rounding_floor": floor, + "finite": bool(torch.isfinite(output).all()), + } + result = {"length": length, "page": page, "variants": {}} + for name, times, error in zip(("legacy_half", "repaired_fp32"), timings, errors): + values = torch.tensor(times) + result["variants"][name] = { + "median_ms": float(values.median()), + "p10_ms": float(values.quantile(0.1)), + "p90_ms": float(values.quantile(0.9)), + "samples_ms": times, + "error": error, + } + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--samples", type=int, default=25) + args = parser.parse_args() + if args.output.exists(): + raise FileExistsError(args.output) + if torch.cuda.get_device_capability() != (7, 0): + raise RuntimeError("requires an idle SM70 GPU") + native = Path(flash_attn_v100_cuda.__file__).resolve() + result = { + "torch": torch.__version__, + "cuda": torch.version.cuda, + "device": torch.cuda.get_device_name(), + "native_path": str(native), + "native_sha256": hashlib.sha256(native.read_bytes()).hexdigest(), + "precision_version": flash_attn_v100_cuda.grouped_e4m3_fp32_precision_version(), + "contract": "B1/q8/H6/Hkv1/D256, same E4M3 KV, graph ABBA; not E2E", + "cases": [], + } + for length in (8192, 65536, 131072, 262144): + row = measure(length, 3296, args.samples) + result["cases"].append(row) + print(json.dumps(row), flush=True) + args.output.write_text(json.dumps(result, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/docs/design/sm70_dflash2_fp32_defaults.md b/docs/design/sm70_dflash2_fp32_defaults.md new file mode 100644 index 0000000000..9ac7c6488f --- /dev/null +++ b/docs/design/sm70_dflash2_fp32_defaults.md @@ -0,0 +1,109 @@ +# DFlash2 E4M3 attention with FP32 state + +The SM70 Flash-V100 `fp8` KV alias now resolves to `fp8_e4m3`. Explicit +`fp8_e5m2` remains available for reproducing older deployments. Explicit +formats, checkpoint-resolved formats and non-SM70 backends are preserved. +This changes the FP8 alias, not the general `auto` cache policy or model weights. + +For E4M3 DFlash2 target verification, the backend no longer selects the legacy +grouped entry that stores normalized attention partials in FP16. Its advertised +E4M3 support only proves byte-format compatibility. The backend routes compatible +single-request q2–8/H6/Hkv1/D256 input through the repaired FP32 entry, using +the metadata builder's live per-query lengths, including zero padding. + +The repaired computation retains compensated QK accumulation, compensated +probabilities, tile-local FP32 PV accumulation, FP32 unnormalized partition +numerators, separate FP32 max/sum, and FP32 combination. KV storage remains +E4M3; Tensor Core operands and the final activation remain FP16. This is not a +full-FP32 model and does not remove FP8 quantization error. + +Precision capability revision **4** adds page sizes **1728/3456** to the existing +800/848/1616/1648/3296 runtime-stride implementation. The arithmetic is unchanged +from repaired revision 3. It also adds FP32 scalar E4M3 partial storage and +reduction. The wrapper requires revision 4 so a new page or FP32 workspace +cannot reach an incompatible native binary. Rebuild the extension and restart +workers together with this Python update. + +Unsupported grouped shapes, independent-request batches, q16 and explicit grouped +rollback use the scalar path with FP32 accumulation and the new FP32 split +workspace. Scalar and XQA workspaces are separated by partial dtype in the cache. +They do not re-enter the legacy E4M3 FP16-partial verifier. If the native binary +lacks revision 4, E4M3 scalar calls raise a rebuild error instead of silently +storing half partials. A DFlash2 target q1 also uses FP32 scalar state instead of +the half-partial XQA wave route. Ordinary non-DFlash XQA dispatch is unchanged. +The low-level legacy operator remains available for numerical A/B tests and +explicit E5M2 compatibility; it is not the E4M3 serving policy. + +The Qwen3.8 DFlash2 configuration also enables +`VLLM_SM70_DFLASH2_FP32_LOGITS=1` by default so candidate rerank and dense fallback +retain FP32 logits. This is model-scoped configuration; the global environment +default remains off for unrelated models. Explicit environment overrides are +preserved. This does not enable MTP or change the sampling distribution settings. + +## Observability and reproduction + +A compatible worker reports `E4M3 grouped FP32 route selected` with its page, +query rows and FP32 state representation. Route summaries include +`prefill_smallq_e4m3_grouped_fp32` and `fp8_kv_decode_grouped_fp32`. +Check these alongside the resolved KV dtype and FP32-logit preparation; an image +name or requested flag is insufficient evidence of the numerical route. + +With an isolated Python 3.12 runtime, CUDA 12.8, Torch 2.10+cu128 and idle V100s: + +```bash +CUDA_VISIBLE_DEVICES=1 .venv/bin/python -m pytest -q \ + tests/v1/attention/test_sm70_flash_v100_policy.py \ + tests/v1/attention/test_sm70_e4m3_grouped.py \ + tests/v1/spec_decode/test_dflash2.py \ + tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py + +CUDA_VISIBLE_DEVICES=1 .venv/bin/python \ + benchmarks/benchmark_sm70_dflash2_fp32_attention.py --output operator.json +``` + +Point `PYTHONPATH` at the task's Python package and rebuilt extension. The +operator benchmark holds E4M3 bytes, query, causal visibility, KV scales and +native library fixed. It reports 20 warmups, 25 ABBA blocks (50 samples per +variant), latency percentiles, PyTorch FP32/FP64 arithmetic error and the FP16 +output-rounding floor. It separates avoidable attention error from cache +quantization, and is not a model-throughput benchmark. + +Model acceptance must be checked with request-level draft/accepted-token +counter deltas. A rolling logger's mean acceptance length does not provide a +matched request comparison. FP32 is the chosen arithmetic contract, but lower +operator L2 alone does not establish model-quality improvement or guarantee +identical sampled tokens. + +## Rebuilt operator results, 2026-09-08 + +The initial routing build passed **420 tests** before the scalar fallback was +extended. The native q8 tests include 8K/64K/128K/256K, newly admitted +1728/3456 pages, relocated pages, non-unit scales and CUDA Graph replay with +changed/zero row lengths. Numerical assertions bound error relative to the +FP16 output-rounding floor, rather than only using an aggregate loose tolerance. +The final scalar/q1 policy follow-up passes **338 checks**, including all five +new scalar tests. The separate native/planner run passed 102 checks and exposed +an incorrect test assertion that eager and graph streams must share a workspace; +the corrected tests check reuse within one stream and FP32 buffers in both. +This was a test expectation error, not an attention numerical failure. + +Physical GPU2, V100-SXM2-32GB, Torch 2.10.0+cu128, CUDA 12.8.93/GCC12, +q8/page3296, one rebuilt library and paired graph samples: + +| Context | Legacy ms | FP32 ms | Legacy relative L2 | FP32 relative L2 | +| --- | ---: | ---: | ---: | ---: | +| 8192 | 0.0778 | 0.1034 | 0.00035316 | 0.00020792 | +| 65536 | 0.2714 | 0.4209 | 0.00034410 | 0.00020785 | +| 131072 | 0.4977 | 0.7916 | 0.00035442 | 0.00020903 | +| 262144 | 0.9462 | 1.5206 | 0.00033226 | 0.00020327 | + +L2 uses PyTorch FP32 attention over identical quantized KV. The companion FP64 +reference gives an FP16 rounding floor of 0.000203267 at 256K. The operator +has about 39% lower L2 and 61% higher latency at that length; this is a deliberate +precision choice, not a speedup. It does not establish a model acceptance gain. +No cross-GPU absolute timing is combined. + +Native library SHA256: +`d2f70b502985af14fe816b379ffa319ca4887191dfab311d62836ee121a41ef3`. +Raw artifacts are indexed under `dflash2-e4m3-fp32-default-20260908`, including +`operator-r2.json` with all samples and both numerical references. diff --git a/docs/design/sm70_e4m3_grouped_fp32.md b/docs/design/sm70_e4m3_grouped_fp32.md index 20e92abe0b..5dd8251a48 100644 --- a/docs/design/sm70_e4m3_grouped_fp32.md +++ b/docs/design/sm70_e4m3_grouped_fp32.md @@ -1,5 +1,11 @@ # Experimental E4M3 grouped attention with FP32 partial state +**2026-09-08 policy update:** [DFlash2 FP32 defaults](sm70_dflash2_fp32_defaults.md) +supersede the default-routing and KV-alias decisions recorded below. E4M3 +DFlash2 verification now uses repaired FP32 state; revision 4 adds the +1728/3456 page layouts. Earlier measurements and failed model gates retain +their original artifact attribution. + ## Scope and admission **Current mainline audit decision (2026-09-07): enabled for compatible diff --git a/docs/design/sm70_glm53_flash_nvfp4.md b/docs/design/sm70_glm53_flash_nvfp4.md index 198357ef66..d0710249eb 100644 --- a/docs/design/sm70_glm53_flash_nvfp4.md +++ b/docs/design/sm70_glm53_flash_nvfp4.md @@ -114,7 +114,9 @@ The adaptation is divided into independently testable surfaces: workspace before two Tensor Core GEMMs. The older direct scalar kernel is retained only as a reference/test path and is not accepted for B1 decode. Use the explicit `fp8_e4m3` cache dtype because the historical generic SM70 - `fp8` alias resolves to E5M2 for other model families. + `fp8` alias historically resolved to E5M2 for other model families. The + [E4M3 default update](sm70_dflash2_fp32_defaults.md) changes that alias; + explicit `fp8_e4m3` keeps this recipe independent of that version boundary. 9. Keep all GLM mHC4/H4096 execution on native SM70 kernels. Small-M fused decode follows the DeepSeek-V4 FP32 staging design, but its final Sinkhorn, residual mix, and RMSNorm stage is a dedicated single-CTA CUDA kernel for diff --git a/docs/design/sm70_v100_migration_control.md b/docs/design/sm70_v100_migration_control.md index 7592d4c328..f114bbc08b 100644 --- a/docs/design/sm70_v100_migration_control.md +++ b/docs/design/sm70_v100_migration_control.md @@ -2,6 +2,19 @@ Date: 2026-05-30 +## DFlash2 E4M3 FP32 default policy, 2026-09-08 + +[The precision-default change](sm70_dflash2_fp32_defaults.md) is based on +`56f534e672657a6c7599afd6c0dcb2e2c211b2e3`. It removes E4M3 admission to the +legacy FP16-partial verifier, admits DFlash2 pages 1728/3456 to repaired FP32 +attention, resolves the SM70 `fp8` alias to E4M3, and enables FP32 logits in +the existing Qwen3.8 DFlash2 default configuration. Explicit E5M2 remains +available. No MTP enablement is part of this change. + +The legacy half-partial microbenchmark is retained for the speed/precision +comparison; its higher speed is not a reason to restore it to E4M3 serving. +Only measured model results may establish acceptance or throughput effects. + ## v37 prefill integration: model-parity hold, 2026-09-07 Draft [PR548](https://github.com/1CatAI/1Cat-vLLM/pull/548) integrates the diff --git a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py index 9c360a3bc0..30cd71208f 100644 --- a/flash-attention-v100/flash_attn_v100/flash_attn_interface.py +++ b/flash-attention-v100/flash_attn_v100/flash_attn_interface.py @@ -124,11 +124,12 @@ def _allocate_decode_workspace( num_heads: int, head_dim: int, max_num_partitions: int, + partial_dtype: torch.dtype = torch.float16, ) -> _DecodeWorkspace: return _DecodeWorkspace( tmp_out=torch.empty( (batch_capacity, num_heads, max_num_partitions, head_dim), - dtype=torch.float16, + dtype=partial_dtype, device=q.device, ), max_logits=torch.empty( @@ -313,6 +314,7 @@ def _get_decode_workspace_for_plan( head_dim: int, plan: _DecodePlan, active_num_partitions: torch.Tensor | None = None, + partial_dtype: torch.dtype = torch.float16, ): device_index = q.device.index if q.device.index is not None else -1 stream_id = _workspace_stream_id(q.device) @@ -323,6 +325,7 @@ def _get_decode_workspace_for_plan( num_heads, head_dim, plan.partition_size, + partial_dtype, ) workspace = _decode_workspace_cache.get(key) if _can_cache_workspace(q) else None @@ -338,6 +341,7 @@ def _get_decode_workspace_for_plan( max_num_partitions=_round_decode_partition_capacity( plan.workspace_num_partitions ), + partial_dtype=partial_dtype, ) if _can_cache_workspace(q): _decode_workspace_cache[key] = workspace @@ -967,6 +971,11 @@ def flash_attn_decode_paged( if softmax_scale is None: softmax_scale = q.shape[-1] ** -0.5 + e4m3_fp32 = kv_cache_dtype in ("fp8", "fp8_e4m3") + if e4m3_fp32 and not flash_attn_grouped_e4m3_fp32_available(): + raise RuntimeError( + "Rebuild Flash-V100 for E4M3 FP32 scalar decode (precision revision 4)" + ) q = maybe_contiguous(q) block_table = maybe_contiguous(block_table) seq_lens = maybe_contiguous(seq_lens) @@ -1011,6 +1020,7 @@ def flash_attn_decode_paged( head_dim=head_dim, plan=plan, active_num_partitions=active_num_partitions, + partial_dtype=torch.float32 if e4m3_fp32 else torch.float16, ) ) @@ -1078,7 +1088,7 @@ def flash_attn_grouped_e4m3_fp32_available() -> bool: return ( hasattr(flash_attn_v100_cuda, "grouped_e4m3_fp32_paged_fwd") and callable(version) - and int(version()) >= 3 + and int(version()) >= 4 ) @@ -1094,7 +1104,7 @@ def flash_attn_grouped_e4m3_fp32_paged( k_scale: float = 1.0, v_scale: float = 1.0, ) -> torch.Tensor: - """Experimental E4M3 q2..8/GQA6/D256 attention over one KV sequence. + """E4M3 q2..8/GQA6/D256 attention over one KV sequence. Row lengths are authoritative GPU metadata, not inferred from padded Q. Zero lengths produce zero outputs. All positive lengths must fit the @@ -1103,10 +1113,11 @@ def flash_attn_grouped_e4m3_fp32_paged( Tensor Core operands and final output remain FP16. KV must encode E4M3. Precision revision 3 retains FP32 numerators and separate max/sum until the final normalization, as well as compensated QK/P and tile-local PV. + Revision 4 adds DFlash2 1728/3456 pages and FP32 scalar fallback workspace. """ if not flash_attn_grouped_e4m3_fp32_available(): raise RuntimeError( - "Rebuild Flash-V100 for E4M3 grouped FP32 precision revision 3" + "Rebuild Flash-V100 for E4M3 grouped FP32 precision revision 4" ) workspace = _get_grouped_verify_workspace(q, partial_dtype=torch.float32) return flash_attn_v100_cuda.grouped_e4m3_fp32_paged_fwd( diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 5e96615420..3d9df0ea23 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -1026,10 +1026,11 @@ __device__ __forceinline__ float dot_qk_cache(const __half* __restrict__ q_ptr, } template + int SEQ_LEN_ROUTE = kXQARouteAllSeqLens, bool ANCHORED_SWA = false, + typename PARTIAL_T = __half> __global__ void flash_attention_decode_partition_kernel( const __half* __restrict__ q, const void* __restrict__ k_cache, - const void* __restrict__ v_cache, __half* __restrict__ tmp_out, + const void* __restrict__ v_cache, PARTIAL_T* __restrict__ tmp_out, float* __restrict__ max_logits, float* __restrict__ exp_sums, const int* __restrict__ block_table, const int* __restrict__ seq_lens, const int* __restrict__ active_num_partitions, const int batch_size, @@ -1210,7 +1211,11 @@ __global__ void flash_attention_decode_partition_kernel( const float out_scale = KV_DTYPE == flash_v100::KV_CACHE_DTYPE_FP16 ? inv_part_sum : inv_part_sum * v_scale; - tmp_out[tmp_out_base + d] = __float2half(acc * out_scale); + if constexpr (std::is_same_v) { + tmp_out[tmp_out_base + d] = acc * out_scale; + } else { + tmp_out[tmp_out_base + d] = __float2half(acc * out_scale); + } } if (threadIdx.x == 0) { @@ -2680,9 +2685,10 @@ __launch_bounds__(kGroupedVerifyThreads) void flash_attention_grouped_verify_e5m } } -template +template __global__ void flash_attention_decode_reduce_kernel( - const __half* __restrict__ tmp_out, const float* __restrict__ max_logits, + const PARTIAL_T* __restrict__ tmp_out, const float* __restrict__ max_logits, const float* __restrict__ exp_sums, const int* __restrict__ seq_lens, const int* __restrict__ active_num_partitions, __half* __restrict__ out, const int batch_size, const int max_num_partitions, const int num_heads_q, @@ -2767,11 +2773,11 @@ __global__ void flash_attention_decode_reduce_kernel( for (int d = threadIdx.x; d < D; d += blockDim.x) { float acc = 0.f; for (int i = 0; i < num_partitions; ++i) { - acc = fmaf( - weight_shared[i], - __half2float(tmp_out[tmp_out_base + - static_cast(i) * tmp_out_stride2 + d]), - acc); + acc = fmaf(weight_shared[i], + static_cast( + tmp_out[tmp_out_base + + static_cast(i) * tmp_out_stride2 + d]), + acc); } out[out_base + d] = __float2half(acc * inv_global_sum); } @@ -3437,7 +3443,7 @@ __global__ void flash_attention_decode_qk_scores_kernel( } template + int SEQ_LEN_ROUTE = kXQARouteAllSeqLens, typename PARTIAL_T = __half> void launch_flash_attention_decode_paged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, at::Tensor& out, const at::Tensor& block_table, const at::Tensor& seq_lens, @@ -3469,11 +3475,11 @@ void launch_flash_attention_decode_paged( const auto launch_partition = [&](auto anchored_tag) { constexpr bool kAnchored = decltype(anchored_tag)::value; flash_attention_decode_partition_kernel + SEQ_LEN_ROUTE, kAnchored, PARTIAL_T> <<>>( reinterpret_cast(q.data_ptr()), k_cache.data_ptr(), v_cache.data_ptr(), - reinterpret_cast<__half*>(tmp_out.data_ptr()), + reinterpret_cast(tmp_out.data_ptr()), max_logits.data_ptr(), exp_sums.data_ptr(), block_table.data_ptr(), seq_lens.data_ptr(), active_num_partitions.data_ptr(), batch_size, max_num_blocks, @@ -3502,9 +3508,10 @@ void launch_flash_attention_decode_paged( return; } - flash_attention_decode_reduce_kernel + flash_attention_decode_reduce_kernel <<>>( - reinterpret_cast(tmp_out.data_ptr()), + reinterpret_cast(tmp_out.data_ptr()), max_logits.data_ptr(), exp_sums.data_ptr(), seq_lens.data_ptr(), active_num_partitions.data_ptr(), reinterpret_cast<__half*>(out.data_ptr()), batch_size, @@ -4279,7 +4286,8 @@ at::Tensor flash_attention_grouped_e4m3_fp32_paged( TORCH_CHECK( k.dim() == 4 && k.size(2) == 1 && k.size(3) == 256 && (k.size(1) == 800 || k.size(1) == 848 || k.size(1) == 1616 || - k.size(1) == 1648 || k.size(1) == 3296) && + k.size(1) == 1648 || k.size(1) == 1728 || k.size(1) == 3296 || + k.size(1) == 3456) && k.scalar_type() == at::kByte && v.scalar_type() == at::kByte && v.sizes() == k.sizes(), "E4M3 grouped FP32 requires supported uint8 paged KV [pages,page,1,256]"); @@ -4366,7 +4374,8 @@ at::Tensor flash_attention_grouped_e4m3_fp32_paged( int64_t flash_attention_grouped_e4m3_fp32_precision_version() { // Revision 3 retains unnormalized FP32 numerators and separate max/sum. // Older normalized-partial/LSE workspaces are not ABI-compatible. - return 3; + // Revision 4 admits DFlash2 1728/3456 pages and FP32 scalar E4M3 partials. + return 4; } int64_t flash_attention_grouped_verify_max_query_tokens() { @@ -4763,7 +4772,10 @@ at::Tensor flash_attention_decode_paged( TORCH_CHECK(k_scale > 0.f && v_scale > 0.f, "fp8 k/v scales must be positive"); } - TORCH_CHECK(tmp_out.dtype() == torch::kFloat16, "tmp_out must be fp16"); + TORCH_CHECK(tmp_out.dtype() == torch::kFloat16 || + (kv_dtype_code == flash_v100::KV_CACHE_DTYPE_FP8_E4M3 && + tmp_out.dtype() == torch::kFloat32), + "tmp_out must be fp16, or fp32 for E4M3 scalar decode"); TORCH_CHECK(max_logits.dtype() == torch::kFloat32, "max_logits must be fp32"); TORCH_CHECK(exp_sums.dtype() == torch::kFloat32, "exp_sums must be fp32"); TORCH_CHECK(block_table.dtype() == torch::kInt32, @@ -4864,21 +4876,31 @@ at::Tensor flash_attention_decode_paged( k_scale, v_scale, window_size_left, window_size_right, stream, 0, 0, 0, \ true, anchor_lens_ptr, static_cast(anchored_window)) -#define LAUNCH_BY_KV_DTYPE(HDIM, PARTITION) \ - do { \ - switch (kv_dtype_code) { \ - case flash_v100::KV_CACHE_DTYPE_FP16: \ - LAUNCH_TYPED(HDIM, PARTITION, flash_v100::KV_CACHE_DTYPE_FP16); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ - LAUNCH_TYPED(HDIM, PARTITION, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ - break; \ - case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ - LAUNCH_TYPED(HDIM, PARTITION, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ - break; \ - default: \ - TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ - } \ +#define LAUNCH_BY_KV_DTYPE(HDIM, PARTITION) \ + do { \ + switch (kv_dtype_code) { \ + case flash_v100::KV_CACHE_DTYPE_FP16: \ + LAUNCH_TYPED(HDIM, PARTITION, flash_v100::KV_CACHE_DTYPE_FP16); \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E4M3: \ + if (tmp_out.scalar_type() == at::kFloat) { \ + launch_flash_attention_decode_paged< \ + HDIM, PARTITION, flash_v100::KV_CACHE_DTYPE_FP8_E4M3, \ + kXQARouteAllSeqLens, float>( \ + q, k_cache, v_cache, out, block_table, seq_lens, tmp_out, \ + max_logits, exp_sums, active_num_partitions, softmax_scale, \ + launch_num_partitions, k_scale, v_scale, window_size_left, \ + window_size_right, stream); \ + } else { \ + LAUNCH_TYPED(HDIM, PARTITION, flash_v100::KV_CACHE_DTYPE_FP8_E4M3); \ + } \ + break; \ + case flash_v100::KV_CACHE_DTYPE_FP8_E5M2: \ + LAUNCH_TYPED(HDIM, PARTITION, flash_v100::KV_CACHE_DTYPE_FP8_E5M2); \ + break; \ + default: \ + TORCH_CHECK(false, "Unsupported kv_cache_dtype: ", kv_cache_dtype); \ + } \ } while (0) #define LAUNCH_BY_PARTITION(HDIM) \ diff --git a/tests/kernels/attention/test_sm70_e4m3_scalar_fp32.py b/tests/kernels/attention/test_sm70_e4m3_scalar_fp32.py new file mode 100644 index 0000000000..a8d8736658 --- /dev/null +++ b/tests/kernels/attention/test_sm70_e4m3_scalar_fp32.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The E4M3 scalar fallback must retain FP32 split state as well.""" + +from types import SimpleNamespace + +import pytest +import torch + + +def test_e4m3_scalar_rejects_stale_native(monkeypatch): + interface = pytest.importorskip("flash_attn_v100.flash_attn_interface") + monkeypatch.setattr(interface, "flash_attn_v100_cuda", SimpleNamespace()) + q = torch.empty((1, 6, 256), dtype=torch.float16) + with pytest.raises(RuntimeError, match="FP32 scalar decode.*revision 4"): + interface.flash_attn_decode_paged(q, q, q, q, q, kv_cache_dtype="fp8_e4m3") + + +@pytest.mark.parametrize( + "rows,dim,page,length,partition", + [ + (1, 64, 16, 2048, 256), + (2, 128, 16, 4096, 512), + (16, 256, 1728, 8192, 1024), + (1, 256, 3456, 262144, 1024), + ], +) +def test_scalar_fp32_workspace_and_live_graph( + monkeypatch, rows, dim, page, length, partition +): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0): + pytest.skip("requires SM70") + interface = pytest.importorskip("flash_attn_v100.flash_attn_interface") + if not interface.flash_attn_grouped_e4m3_fp32_available(): + pytest.skip("rebuild FP32 E4M3 decode") + torch.manual_seed(20260908) + pages = (length + page - 1) // page + q = torch.randn((rows, 6, dim), device="cuda", dtype=torch.float16) + kv = torch.randn((2, pages, page, 1, dim), device="cuda", dtype=torch.float16) + kv[1].add_(3.0) + kv = kv.to(torch.float8_e4m3fn).view(torch.uint8) + k, v = kv.unbind(0) + table = torch.arange(pages, device="cuda", dtype=torch.int32)[None].repeat(rows, 1) + lengths = torch.arange( + length - rows + 1, length + 1, device="cuda", dtype=torch.int32 + ) + original = lengths.clone() + output = torch.empty_like(q) + get_workspace = interface._get_decode_workspace_for_plan + workspaces = [] + + def checked_workspace(*args, **kwargs): + result = get_workspace(*args, **kwargs) + assert result[0].dtype == torch.float32 + workspaces.append(result[0]) + return result + + monkeypatch.setattr(interface, "_get_decode_workspace_for_plan", checked_workspace) + + def call(): + interface.flash_attn_decode_paged( + q, + k, + v, + table, + lengths, + out=output, + kv_cache_dtype="fp8_e4m3", + softmax_scale=dim**-0.5, + k_scale=0.5, + v_scale=1.25, + max_seq_len_hint=length, + workspace_seq_capacity_hint=pages * page, + partition_size_hint=partition, + ) + + call() + call() + assert workspaces[0].data_ptr() == workspaces[1].data_ptr() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + call() + rk = k.reshape(-1, dim)[:length].view(torch.float8_e4m3fn).double() * 0.5 + rv = v.reshape(-1, dim)[:length].view(torch.float8_e4m3fn).double() * 1.25 + for zero in (False, True, False): + lengths.copy_(original) + if zero: + lengths[-1] = 0 + graph.replay() + scores = q.transpose(0, 1).double() @ rk.T * dim**-0.5 + scores.masked_fill_( + torch.arange(length, device="cuda")[None, None] >= lengths[None, :, None], + -torch.inf, + ) + expected = (scores.softmax(-1).nan_to_num(0) @ rv).transpose(0, 1) + assert bool(torch.isfinite(output).all()) + assert bool((output[lengths == 0] == 0).all()) + if bool((lengths != 0).any()): + error = (output.double() - expected).norm() / expected.norm() + floor = (expected.half().double() - expected).norm() / expected.norm() + assert float(error) <= 1.04 * float(floor) + 2e-6 diff --git a/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py b/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py index 6fc70ab8b6..338b8c79e7 100644 --- a/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py +++ b/tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py @@ -18,7 +18,7 @@ def _native(): @pytest.mark.parametrize("has_entry", [False, True]) -@pytest.mark.parametrize("version", [None, 0, 1, 2, 3, 4]) +@pytest.mark.parametrize("version", [None, 0, 1, 2, 3, 4, 5]) def test_precision_capability_rejects_stale_binary(monkeypatch, has_entry, version): interface = pytest.importorskip("flash_attn_v100.flash_attn_interface") native = SimpleNamespace() @@ -28,7 +28,7 @@ def test_precision_capability_rejects_stale_binary(monkeypatch, has_entry, versi native.grouped_e4m3_fp32_precision_version = lambda: version monkeypatch.setattr(interface, "flash_attn_v100_cuda", native) assert interface.flash_attn_grouped_e4m3_fp32_available() is ( - has_entry and version is not None and version >= 3 + has_entry and version is not None and version >= 4 ) @@ -40,6 +40,13 @@ def test_precision_capability_rejects_stale_binary(monkeypatch, has_entry, versi (8, 1616, 65536), (5, 1648, 131072), (5, 3296, 262144), + # DFlash2 q8 uses the same repaired arithmetic at each context boundary. + (8, 3296, 8192), + (8, 3296, 65536), + (8, 3296, 131072), + (8, 3296, 262144), + (8, 1728, 131072), + (8, 3456, 262144), ], ) def test_fp32_grouped_row_lengths_graph(rows, page, length): diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 90f4aef6f7..2525878271 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -301,7 +301,7 @@ def test_sm70_e5m2_decode_fast_route_envs_are_default_on(monkeypatch): assert envs.VLLM_FLASH_V100_XQA_E5M2_P1024_BEGIN == 49152 -def test_sm70_flash_v100_fp8_alias_resolves_to_e5m2(monkeypatch): +def test_sm70_flash_v100_fp8_alias_resolves_to_e4m3(monkeypatch): import vllm.engine.arg_utils as arg_utils import vllm.envs as envs @@ -318,6 +318,10 @@ def test_sm70_flash_v100_fp8_alias_resolves_to_e5m2(monkeypatch): assert ( arg_utils._resolve_sm70_flash_v100_kv_cache_dtype_alias("fp8", "fp8") + == "fp8_e4m3" + ) + assert ( + arg_utils._resolve_sm70_flash_v100_kv_cache_dtype_alias("fp8_e5m2", "fp8_e5m2") == "fp8_e5m2" ) assert ( @@ -1810,15 +1814,14 @@ def grouped_verify( num_query_tokens=query_len, ) grouped_verify.supports_e4m3 = True # type: ignore[attr-defined] - if query_len == 16: - assert not impl._dflash2_grouped_verify_allowed( - query, - key_cache, - value_cache, - attn_metadata, - num_query_tokens=query_len, - ) - return + assert not impl._dflash2_grouped_verify_allowed( + query, + key_cache, + value_cache, + attn_metadata, + num_query_tokens=query_len, + ) + return result = impl._flash_v100_small_query_prefill_as_decode( layer, query, @@ -1841,6 +1844,84 @@ def grouped_verify( assert torch.all(output == 1) +@pytest.mark.parametrize("page", [1648, 1728, 3296, 3456]) +@pytest.mark.parametrize("native_available", [False, True]) +def test_dflash2_e4m3_cannot_bypass_fp32_with_legacy_verifier(page, native_available): + from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl + + impl = FlashAttnV100Impl( + num_heads=6, + head_size=256, + scale=0.0625, + num_kv_heads=1, + alibi_slopes=None, + sliding_window=None, + kv_cache_dtype="fp8_e4m3", + ) + impl.use_dflash2_grouped_verify = True + impl.use_smallq_decode_xqa = True + impl.dflash2_grouped_verify_max_query_tokens = 16 + calls = [] + + def legacy(*args, **kwargs): + pytest.fail("E4M3 DFlash2 selected FP16 partial state") + + legacy.supports_e4m3 = True # type: ignore[attr-defined] + + def precise(q, k, v, table, lengths, **kwargs): + calls.append(("fp32", table, lengths)) + assert kwargs["k_scale"] == 0.5 + assert kwargs["v_scale"] == 1.25 + kwargs["out"].fill_(3) + + def scalar(q, k, v, table, lengths, **kwargs): + calls.append(("scalar", table, lengths)) + kwargs["out"].fill_(2) + + impl.flash_attn_grouped_verify_paged = legacy + impl.flash_attn_grouped_e4m3_fp32_paged = precise if native_available else None + impl._call_flash_attn_decode_paged = scalar + q = torch.zeros((8, 6, 256), dtype=torch.float16) + kv = torch.zeros((2, page, 1, 256), dtype=torch.uint8) + out = torch.empty_like(q) + table = torch.tensor([[1, 0]], dtype=torch.int32) + row_table = table.repeat(8, 1) + # Padding must retain zero visibility rather than shift the causal boundary. + lengths = torch.tensor( + [2049, 2050, 2051, 2052, 2053, 2054, 0, 0], dtype=torch.int32 + ) + metadata = SimpleNamespace( + num_actual_tokens=8, + causal=True, + is_dflash_selector_target=True, + max_model_len=262144, + block_table=table, + seq_lens=torch.tensor([2054], dtype=torch.int32), + query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + smallq_decode_block_table=row_table, + smallq_decode_seq_lens=lengths, + smallq_query_start_loc=torch.tensor([0, 8], dtype=torch.int32), + ) + result = impl._flash_v100_small_query_prefill_as_decode( + SimpleNamespace(_k_scale_float=0.5, _v_scale_float=1.25), + q, + kv, + kv, + metadata, + out, + metadata.query_start_loc, + metadata.seq_lens, + ) + assert result is out + assert len(calls) == 1 + assert calls[0][0] == ("fp32" if native_available else "scalar") + assert ( + calls[0][1].data_ptr() == (table if native_available else row_table).data_ptr() + ) + assert calls[0][2].data_ptr() == lengths.data_ptr() + assert bool((out == (3 if native_available else 2)).all()) + + def test_flash_v100_batched_grouped_workspace_preserves_single_request_layout(): from flash_attn_v100.flash_attn_interface import _get_grouped_verify_workspace @@ -2154,7 +2235,8 @@ def fail_scalar(*args, **kwargs): assert torch.all(output == 1) -def test_flash_v100_decode_uses_xqa_for_e4m3_g6_d256(monkeypatch): +@pytest.mark.parametrize("dflash_target", [False, True]) +def test_flash_v100_decode_e4m3_respects_dflash_fp32_policy(monkeypatch, dflash_target): from vllm.v1.attention.backends.flash_attn_v100 import FlashAttnV100Impl monkeypatch.delenv("VLLM_FLASH_V100_DECODE_USE_XQA", raising=False) @@ -2181,13 +2263,16 @@ def hit_xqa(*args, **kwargs): ) kwargs["out"].fill_(1) - def fail_scalar(*args, **kwargs): - raise AssertionError("E4M3 G6/D256 decode should select XQA") + def hit_scalar(*args, **kwargs): + assert dflash_target + calls.append(("scalar", None, kwargs.get("kv_cache_dtype"))) + kwargs["out"].fill_(1) impl.flash_attn_decode_paged_xqa = hit_xqa # type: ignore[method-assign] - impl.flash_attn_decode_paged = fail_scalar # type: ignore[method-assign] + impl.flash_attn_decode_paged = hit_scalar # type: ignore[method-assign] attn_metadata = SimpleNamespace( num_actual_tokens=1, + is_dflash_selector_target=dflash_target, block_table=torch.tensor([[0]], dtype=torch.int32), seq_lens=torch.tensor([1025], dtype=torch.int32), flash_v100_decode_max_seq_len_hint=1025, @@ -2210,7 +2295,9 @@ def fail_scalar(*args, **kwargs): ) assert result is output - assert calls == [("xqa", 64, "fp8_e4m3")] + assert calls == [ + ("scalar", None, "fp8_e4m3") if dflash_target else ("xqa", 64, "fp8_e4m3") + ] assert torch.all(output == 1) @@ -3198,7 +3285,11 @@ def scalar(*args, **kwargs): assert calls[0][1]["v_scale"] == 1.25 if mode == "selected": assert calls[0][0][3] is metadata.block_table - assert routes == ["prefill_smallq_e4m3_grouped_fp32"] + assert routes == [ + "fp8_kv_decode", + "fp8_kv_decode_grouped_fp32", + "prefill_smallq_e4m3_grouped_fp32", + ] assert bool((out == 3).all()) else: assert calls[0][0][3] is table diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 62c3f7a9c1..4809b91f21 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -85,6 +85,8 @@ _SM70_SPECULATIVE_AUX_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 9, 18) _SM70_DFLASH2_VERIFIER_DEFAULTS = { + # Preserve candidate and dense logits in FP32 through sampling. + "VLLM_SM70_DFLASH2_FP32_LOGITS": "1", # This is the target projection's memory-neutral FP8 layout, not the # rejected draft-MLP QPN8 experiment. Per-layer TP/shape checks retain the # original layout whenever the exact operator contract is unavailable. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 6b73e33f00..ac032debed 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -142,7 +142,7 @@ def _resolve_sm70_flash_v100_kv_cache_dtype_alias( requested_dtype: str, resolved_dtype: str, ) -> str: - """Preserve the historical 1Cat V100 meaning of the ``fp8`` alias.""" + """Use E4M3 storage for the SM70 Flash-V100 ``fp8`` alias.""" if ( requested_dtype != "fp8" or resolved_dtype != "fp8" @@ -157,12 +157,11 @@ def _resolve_sm70_flash_v100_kv_cache_dtype_alias( if capability is None or (capability.major, capability.minor) != (7, 0): return resolved_dtype logger.warning_once( - "On SM70 Flash-V100, --kv-cache-dtype fp8 resolves to fp8_e5m2 " - "for compatibility with the optimized 1Cat V100 KV-cache path. " - "Use explicit fp8_e4m3 to request E4M3 KV storage. Model weight " + "On SM70 Flash-V100, --kv-cache-dtype fp8 resolves to fp8_e4m3. " + "Use explicit fp8_e5m2 to request legacy E5M2 KV storage. Model weight " "quantization is configured independently." ) - return "fp8_e5m2" + return "fp8_e4m3" # object is used to allow for special typing forms diff --git a/vllm/v1/attention/backends/flash_attn_v100.py b/vllm/v1/attention/backends/flash_attn_v100.py index bd7ee3becc..5496a2ce94 100644 --- a/vllm/v1/attention/backends/flash_attn_v100.py +++ b/vllm/v1/attention/backends/flash_attn_v100.py @@ -4455,9 +4455,9 @@ def __init__(self, *args, **kwargs): ) if use_e4m3_fp32 and self.flash_attn_grouped_e4m3_fp32_paged is None: logger.warning_once( - "E4M3 grouped FP32 requires Flash-V100 precision revision 3; " - "using the existing attention fallback. Rebuild the extension " - "and restart workers to enable the repaired route.", + "E4M3 grouped FP32 requires Flash-V100 precision revision 4; " + "the E4M3 scalar fallback also requires this revision for " + "FP32 partial storage. Rebuild the extension and restart workers.", scope="process", ) self.dflash2_grouped_verify_max_query_tokens = ( @@ -5413,20 +5413,10 @@ def _dflash2_grouped_verify_allowed( and value_cache.dtype == torch.uint8 and key_cache.stride(-1) == 1 and value_cache.stride(-1) == 1 - and ( - self.kv_cache_dtype == "fp8_e5m2" - or ( - self.kv_cache_dtype == "fp8_e4m3" - and num_query_tokens == 8 - and key_cache.stride(0) % 16 == 0 - and key_cache.stride(1) % 16 == 0 - and value_cache.stride(0) % 16 == 0 - and value_cache.stride(1) % 16 == 0 - and getattr( - self.flash_attn_grouped_verify_paged, "supports_e4m3", False - ) - ) - ) + # This legacy verifier stores normalized partials in FP16. + # E4M3 must reach the repaired FP32 path below, including when + # the old native entry advertises E4M3 byte-format support. + and self.kv_cache_dtype == "fp8_e5m2" and block_table is not None and block_table.ndim == 2 and block_table.shape[0] == num_reqs @@ -5612,12 +5602,13 @@ def _call_flash_attn_smallq_decode_paged( v_scale=float(layer._v_scale_float), ) logger.info_once( - "FLASH_ATTN_V100 experimental E4M3 grouped FP32 route " - "selected (rows=%d, page=%d, explicit row lengths).", + "FLASH_ATTN_V100 E4M3 grouped FP32 route selected " + "(rows=%d, page=%d, FP32 numerator/max/sum, explicit row lengths).", query.shape[0], key_cache.shape[1], scope="process", ) + _log_fp8_kv_cache_route("decode", self.kv_cache_dtype, "grouped_fp32") _record_route("prefill_smallq_e4m3_grouped_fp32") return window_size = self._flash_v100_window_size(causal=True) @@ -6816,6 +6807,9 @@ def _flash_v100_decode( self.kv_cache_dtype in ("fp8", "fp8_e4m3") and key_cache.dtype == torch.uint8 and value_cache.dtype == torch.uint8 + # The E4M3 XQA wave route retains half partials. A DFlash2 + # target q1 must honor the same FP32 state policy as q8. + and not getattr(attn_metadata, "is_dflash_selector_target", False) ) ) diff --git a/vllm/v1/attention/ops/sm70_e4m3_grouped.py b/vllm/v1/attention/ops/sm70_e4m3_grouped.py index d113e50635..b4dea81606 100644 --- a/vllm/v1/attention/ops/sm70_e4m3_grouped.py +++ b/vllm/v1/attention/ops/sm70_e4m3_grouped.py @@ -43,7 +43,7 @@ def grouped_e4m3_fp32_allowed( and out.device == query.device and out.is_contiguous() and k.ndim == 4 - and k.shape[1] in (800, 848, 1616, 1648, 3296) + and k.shape[1] in (800, 848, 1616, 1648, 1728, 3296, 3456) and k.shape[2:] == (1, 256) and v.shape == k.shape and k.dtype == torch.uint8