From 4050afe66283afd38a6a5baee7c8ae8572b2c3eb Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 12 Aug 2026 22:28:41 +0700 Subject: [PATCH] Match Qwen SDPA dispatch --- PARITY.md | 23 +++++++++++++++++ src/h3_blackwell_runtime/qwen3vl_text.py | 32 ++++++++++++++++++------ 2 files changed, 47 insertions(+), 8 deletions(-) diff --git a/PARITY.md b/PARITY.md index 76bec82..a351d11 100644 --- a/PARITY.md +++ b/PARITY.md @@ -41,15 +41,38 @@ diagnose preview output. | Qwen layer-0 normalization | Direct `F.rms_norm` versus `qwen0_norm1.pt` | Mean absolute delta about `8.96e-06` | Passed to expected BF16-level tolerance | | Qwen layer-0 attention, before projection fix | Direct versus `qwen0_attention.pt` | Mean absolute delta about `0.0357`; layer output about `0.1496` | First material Qwen divergence was at projection/attention boundary, not tokenization, embeddings, RMSNorm, or RoPE | | Qwen attention backend | Comfy source inspection | Qwen chooses Comfy small-input attention, i.e. SDPA path, not Sage2 | Direct Qwen was changed to BF16 causal-mask SDPA with GQA | +| Qwen SDPA dispatch audit and fix | Actual Spark Comfy block audit found direct raw SDPA omitted Comfy's four-dimensional mask wrapper, backend priority, and masked-GQA fallback. Direct ported that policy without importing Comfy. | Replaying immutable layer-0 capture after the change: attention mean/max `0.000244372` / `0.0161562`; post-attention `0.00026845` / `0.0213833`; layer output `0.000808998` / `0.0526199`. Previous attention mean was about `0.0357`; previous layer-output mean about `0.1496`. | Active code mismatch identified and substantially repaired. The remaining layer output error is now downstream of projection/elementwise BF16 boundaries, not an unreviewed SDPA dispatcher difference. | | Qwen NVFP4 metadata | Qwen layer-0 checkpoint sidecars inspected | Every Q/K/V/O and MLP projection has `full_precision_matrix_mult: true`; `o_proj` and `down_proj` also have `pre_quant_scale` | Direct was incorrectly quantizing activations for all Qwen projections | | Qwen NVFP4 fix | Commit `0cad1db` | Direct honors `full_precision_matrix_mult` and loads/applies `pre_quant_scale`; Spark smoke passed finite expected-shape Q and O projections | Structural fix passed; post-fix numeric QKV/MLP/layer-50 comparisons are still required | | Qwen layer-0 projections, after NVFP4 fix | Fresh Comfy projection capture at `/tmp/fl2va-qwen0-projections`; direct compared from captured input embedding and post-attention state | Q mean/max `6.6547e-05` / `0.00195485`; K `9.11704e-05` / `0.00241077`; V `6.0298e-05` / `0.00102112`; gate `0.000365695` / `0.0082469`; up `0.000335494` / `0.00792789`; activated `4.76882e-05` / `0.031621`; down `0.000557594` / `0.0749016` | The full-precision metadata fix removed the previous large projection failure. Projection parity is close but not strict; `down_proj` is the largest remaining layer-0 projection boundary. Preserve these numbers and do not recapture this case. | +| Qwen projection code audit, post-fix | Statement-level comparison against actual Spark Comfy commit `43cb4fffc89bba20ab7bd61467a36d0339338dab`, Torch `2.9.1+cu130`, and `/h3-models/text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors` | All 350 Qwen projection sidecars set `full_precision_matrix_mult: true`; exactly 100 (`o_proj` and `down_proj` in each layer) have BF16 `pre_quant_scale`; no Qwen projection has bias. Packed weight is contiguous uint8, block scale is native E4M3, tensor scale is scalar FP32. | Direct and Comfy agree on the active full-precision branch, NVFP4 scale association, logical shapes, no-bias behavior, pre-scale order, BF16 compute dtype, and use the same installed Comfy Kitchen NVFP4 dequantizer. Do not change projection math based on speculation. | | H3 DiT block trace | `/tmp/fl2va-full-capture` | Block 0 was close, around mean absolute delta `0.00018`; small independent differences accumulated across blocks | Full 50-block strict parity remains failed/unresolved | | Sampler initial state | Comfy `initial.pt` inspected against direct preview | Comfy carries joint AV state; direct preview initializes video noise only and leaves audio zero | Known direct mismatch; identical integer seed is not sampler parity | | Preview output, before Qwen fixes | Dragon/locomotive prompt previews | Flower-like output | Confirmed incorrect conditioning had visible semantic impact | | Preview output, after Qwen metadata fix | 72-frame cat, 12-step direct preview | Cat semantic structure appears, with persistent 32-pixel-like square artifacts | Conditioning improved; denoiser/sampler/VAE boundary remains non-parity | | Video encoder | Direct ffmpeg raw RGB to H.264/yuv420p, no audio | Different from Comfy `CreateVideo`/`SaveVideo` | Not a valid latent-quality gate; cannot explain model-space corruption by itself | +## Qwen Projection Code Audit: Remaining Non-Active Differences + +These are real implementation differences, but the audit established they are +not active numerical branches for the captured Qwen checkpoint. They must not +be presented as an explanation for the current `down_proj` delta without new +evidence. + +| Direct location | Difference from Comfy | Active for current capture? | +| --- | --- | --- | +| `nvfp4.py:47-51` | Calls `.contiguous()` while Comfy preserves loaded storage/view identity. In particular Comfy treats E4M3 block scales as a dtype view, not a numeric cast. | No: checkpoint tensors are already contiguous and direct also uses a dtype view. | +| `nvfp4.py:75-76` | Always flattens and contiguates activation rows; Comfy full-precision `F.linear` retains N-D input/stride. | No known value effect: elements and BF16 dtype are unchanged. | +| `nvfp4.py:82` | Uses `packed_weight.dequantize().to(flat_x)` rather than an explicitly named Comfy compute-dtype dequantization call. | No: input and output compute dtype are BF16, and both call the same installed Kitchen layout/dequantizer. | +| `nvfp4.py:46,81` | Does not implement Comfy global full-precision overrides, disabled-format controls, forced casts, weight functions, or LoRA/offload dispatch. | No: every Qwen sidecar itself forces full precision, no projection has bias, and no patches/LoRAs are attached. | + +**Audit conclusion:** projection source code contains cleanup/generalization gaps, +but no identified active branch mismatch explains the current captured values. +The next code audit must move outward to the exact Qwen block boundary: +Comfy input dtype, full-precision weight construction, SDPA backend selection, +and residual/BF16 rounding boundaries. Do not run another full inference merely +to rediscover this projection result. + ## Current Runtime Scope | Component | Implemented | Known limitation | diff --git a/src/h3_blackwell_runtime/qwen3vl_text.py b/src/h3_blackwell_runtime/qwen3vl_text.py index 37ae283..d060d65 100644 --- a/src/h3_blackwell_runtime/qwen3vl_text.py +++ b/src/h3_blackwell_runtime/qwen3vl_text.py @@ -66,6 +66,29 @@ def _rope(query: torch.Tensor, key: torch.Tensor, theta: float) -> tuple[torch.T return query_output.to(query.dtype), key_output.to(key.dtype) +def _qwen_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Match Comfy's small-input SDPA wrapper for masked Qwen GQA.""" + kwargs = {"attn_mask": mask.unsqueeze(0).unsqueeze(0), "dropout_p": 0.0, "is_causal": False, "enable_gqa": True} + if query.numel() >= 1024 * 128: + from torch.nn.attention import SDPBackend, sdpa_kernel + + priority = [SDPBackend.FLASH_ATTENTION, SDPBackend.CUDNN_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH] + params = torch.backends.cuda.SDPAParams(query, key, value, kwargs["attn_mask"], 0.0, False, True) + supports_native_gqa = ( + torch.backends.cuda.can_use_flash_attention(params) + or torch.backends.cuda.can_use_cudnn_attention(params) + or torch.backends.cuda.can_use_efficient_attention(params) + ) + if not supports_native_gqa: + repeats = query.shape[-3] // key.shape[-3] + key = key.repeat_interleave(repeats, dim=-3) + value = value.repeat_interleave(repeats, dim=-3) + kwargs["enable_gqa"] = False + with sdpa_kernel(priority, set_priority=True): + return F.scaled_dot_product_attention(query, key, value, **kwargs) + return F.scaled_dot_product_attention(query, key, value, **kwargs) + + class _Qwen3VLBlock(nn.Module): def __init__(self, checkpoint: H3Checkpoint, prefix: str, config: Qwen3VL32BTextConfig, dtype: torch.dtype): super().__init__() @@ -99,14 +122,7 @@ class _Qwen3VLBlock(nn.Module): dtype=query.dtype, device=query.device, ).triu_(1) - attention = F.scaled_dot_product_attention( - query, - key, - value, - attn_mask=causal_mask, - is_causal=False, - enable_gqa=True, - ) + attention = _qwen_attention(query, key, value, causal_mask) hidden_states = residual + self.o_proj(attention.transpose(1, 2).reshape(batch, sequence, -1)) residual = hidden_states x = self.post_attention_layernorm(hidden_states)