[Kernel] Add typed SAGE for TRTLLM_ATTN (#5509)

Signed-off-by: lishunyang12 <lishunyang12@163.com>

docs/user_guide/diffusion/attention_backends.md

@@ -14,7 +14,7 @@ The full set of backends and their platform defaults is in the **Backend Options
 
 | Value | Notes |
 |---|---|
-| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
+| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention and **SAGE** quantized attention — see below. Requires `flashinfer`. |
 | `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
 | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
 | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
@@ -167,6 +167,48 @@ sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `T
 Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
 `TRTLLM_ATTN` raises rather than silently degrading.
 
+## TRTLLM_ATTN SAGE Quantization
+
+On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
+FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
+in low precision. Enable it through the typed `quant` block on the attention spec:
+
+| Key | Valid values | Meaning |
+|---|---|---|
+| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
+| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
+| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
+
+V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
+neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
+
+The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
+fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
+`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
+`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
+
+```bash
+vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+  --diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
+      "quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
+```
+
+```python
+from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
+
+AttentionConfig(
+    default=AttentionSpec(
+        backend="TRTLLM_ATTN",
+        quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
+    ),
+)
+```
+
+
+**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
+availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
+(B300); `int8` QK kernels are compiled for **SM100 only**.
+
 ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
 
 Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
@@ -180,22 +222,6 @@ Same prompt and seed across runs. `Total generation time` from `text_to_video.py
 
 Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
 
-## Known Limitations
-
-### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
-
-LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
-
-**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
-
-```bash
-DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Lightricks/LTX-2 ...
-```
-
-### FA4 not yet integrated
-
-FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
 
 ## Choosing a Backend Manually
 
@@ -353,64 +379,3 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_t
     --tensor-parallel-size 2 \
     --output outputs/hv15_sage3.mp4
 ```
-
-### Mixed backends across roles
-
-Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-    --prompt "A dog running across a field of golden wheat." \
-    --diffusion-attention-config.per_role.self.backend FLASH_ATTN \
-    --diffusion-attention-config.per_role.cross.backend TORCH_SDPA \
-    --tensor-parallel-size 2 \
-    --output outputs/wan22_mixed.mp4
-```
-
-### Compare against FlashAttention
-
-Unset the backend override, or explicitly use `FLASH_ATTN`:
-
-```bash
-python examples/offline_inference/text_to_video/text_to_video.py \
-    --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-    --prompt "A dog running across a field of golden wheat." \
-    --height 704 --width 1280 --num-frames 49 \
-    --num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
-    --tensor-parallel-size 2 \
-    --output outputs/wan22_fa3.mp4
-```
-
-## Validation Guidance
-
-Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
-
-Always compare:
-
-- End-to-end runtime
-- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
-- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
-
-At minimum, keep the same:
-
-- model
-- prompt
-- seed
-- resolution
-- frame count / step count
-- parallel config (TP / CFG-parallel / Ulysses degrees)
-
-## Reproducing the Benchmark Table
-
-The end-to-end numbers above were collected by running `text_to_video.py` /
-`text_to_image.py` with the same prompt and seed while varying
-`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
-backends without loading a model:
-
-```bash
-python benchmarks/diffusion/bench_attention_backends.py --preset hv15
-```
-
-It runs all three BF16 backends on representative DiT attention shapes and
-prints a ranking table at the end.

tests/diffusion/attention/test_attention_config.py

@@ -52,6 +52,44 @@ class TestAttentionSpec:
         with pytest.raises(ValueError, match="only supported by the TRTLLM_ATTN"):
             AttentionSpec(backend="TORCH_SDPA", skip_softmax={"target_sparsity": 0.5})
 
+    def test_quant_serialized_with_defaults_and_overrides(self):
+        assert AttentionSpec(backend="TRTLLM_ATTN", quant={"dtype_qk": "fp8_e4m3"}).backend_kwargs()["quant"] == {
+            "dtype_qk": "fp8_e4m3",
+            "q_block_size": 1,
+            "k_block_size": 16,
+        }
+        assert AttentionSpec(
+            backend="TRTLLM_ATTN", quant={"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}
+        ).backend_kwargs()["quant"] == {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}
+
+    def test_quant_superset_fields_passed_through(self):
+        # FLASHINFER_ATTN-style config (dtype_vo / flashinfer_backend) round-trips on the shared spec.
+        bk = AttentionSpec(
+            backend="FLASHINFER_ATTN",
+            quant={"dtype_qk": "bfloat16", "dtype_vo": "fp8_e4m3", "flashinfer_backend": "trtllm-gen"},
+        ).backend_kwargs()["quant"]
+        assert bk["dtype_qk"] == "bfloat16"
+        assert bk["dtype_vo"] == "fp8_e4m3"
+        assert bk["flashinfer_backend"] == "trtllm-gen"
+
+    def test_quant_and_skip_softmax_coexist(self):
+        bk = AttentionSpec(
+            backend="TRTLLM_ATTN", quant={"dtype_qk": "int8"}, skip_softmax={"target_sparsity": 0.5}
+        ).backend_kwargs()
+        assert bk["target_sparsity"] == 0.5 and bk["quant"]["dtype_qk"] == "int8"
+
+    @pytest.mark.parametrize(
+        "spec, match",
+        [
+            ({"backend": "TORCH_SDPA", "quant": {"dtype_qk": "int8"}}, "only supported by the TRTLLM_ATTN"),
+            ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int4"}}, "quant.dtype_qk"),
+            ({"backend": "TRTLLM_ATTN", "quant": {"dtype_qk": "int8", "k_block_size": 8}}, "quant.k_block_size"),
+        ],
+    )
+    def test_quant_validation_rejects(self, spec, match):
+        with pytest.raises(ValueError, match=match):
+            AttentionSpec(**spec)
+
 
 class TestAttentionConfig:
     def test_empty_config(self):

tests/diffusion/attention/test_trtllm_attn.py

@@ -47,6 +47,44 @@ def test_skip_config_pure_resolution():
     assert cfg.resolve_factor(4096, timestep=None) == pytest.approx(0.01 * 4096)
 
 
+def test_quant_config_from_backend_kwargs():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
+
+    assert not QuantConfig.from_backend_kwargs(None).enabled
+    assert not QuantConfig.from_backend_kwargs({}).enabled
+    assert QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "fp8_e4m3"}}) == QuantConfig("fp8_e4m3", 1, 16)
+    qc = QuantConfig.from_backend_kwargs({"quant": {"dtype_qk": "int8", "q_block_size": 4, "k_block_size": 16}})
+    assert qc.enabled and (qc.dtype_qk, qc.q_block_size, qc.k_block_size) == ("int8", 4, 16)
+
+
+def test_quant_rejects_non_sage_dtype():
+    with pytest.raises(RuntimeError, match="supports dtype_qk"):
+        _impl(quant={"dtype_qk": "bfloat16"})
+
+
+def test_quant_quantize_requires_flashinfer_routine(monkeypatch):
+    import vllm_omni.diffusion.attention.backends.trtllm_attn as mod
+
+    monkeypatch.setattr(mod, "_sage_kernel_available", lambda: True)
+    monkeypatch.setattr(mod, "_sage_quantize_fn", lambda: None)
+    with pytest.raises(RuntimeError, match="trtllm_sage_attention_quantize"):
+        _impl(quant={"dtype_qk": "int8"})
+
+
+def test_quant_quantize_calls_routine_and_shapes_sfs():
+    from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
+
+    captured = {}
+
+    def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype):
+        captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype)
+        return "qq", "kq", "vq", "qsfs", "ksfs", "vsfs"
+
+    q_q, k_q, v_q, sfs, blk = QuantConfig(dtype_qk="int8").quantize(object(), object(), object(), fake_quantize)
+    assert (q_q, k_q, v_q, sfs, blk) == ("qq", "kq", "vq", ("qsfs", "ksfs", None, "vsfs"), (1, 16, 0, 1))
+    assert captured == {"q_block_size": 1, "k_block_size": 16, "qk_quant_dtype": torch.int8}
+
+
 def test_skip_factor_none_without_curve():
     assert _impl(target_sparsity=0.5)._resolve_skip_factor(4096) is None
 
@@ -257,6 +295,32 @@ def test_bf16_dense_matches_sdpa():
     assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
 
 
+requires_sage = pytest.mark.skipif(
+    not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None),
+    reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)",
+)
+
+
+@requires_sage
+@pytest.mark.parametrize("dtype_qk", ["fp8_e4m3", "int8"])
+def test_sage_quant_matches_sdpa(dtype_qk):
+    torch.manual_seed(0)
+    B, S, H, D = 2, 256, 8, 128
+    scale = 1.0 / math.sqrt(D)
+    q, k, v = (torch.randn(B, S, H, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+    impl = _impl(quant={"dtype_qk": dtype_qk, "q_block_size": 1, "k_block_size": 16})
+    try:
+        out = impl.forward_cuda(q, k, v, None)
+    except RuntimeError as e:  # pragma: no cover - arch-dependent (int8 is SM100-only)
+        if "Missing TRTLLM-GEN kernel" in str(e):
+            pytest.skip(f"no trtllm-gen SAGE kernel for dtype_qk={dtype_qk} on this GPU arch")
+        raise
+    assert out.shape == (B, S, H, D) and torch.isfinite(out).all()
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out.float() - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high"
+
+
 @requires_trtllm_attn
 def test_skip_tiny_threshold_matches_dense():
     torch.manual_seed(0)

vllm_omni/diffusion/attention/backends/trtllm_attn.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import functools
+import inspect
 import math
 from dataclasses import dataclass, replace
 
@@ -83,6 +85,66 @@ except Exception as e:  # pragma: no cover - import guard
     )
 
 
+@functools.lru_cache(maxsize=1)
+def _sage_kernel_available() -> bool:
+    if not HAS_FLASHINFER:
+        return False
+    try:
+        return "sage_attn_sfs" in inspect.signature(trtllm_ragged_attention_deepseek).parameters
+    except (TypeError, ValueError):
+        return False
+
+
+@functools.lru_cache(maxsize=1)
+def _sage_quantize_fn():
+    try:
+        from flashinfer import trtllm_sage_attention_quantize
+
+        return trtllm_sage_attention_quantize
+    except Exception:  # pragma: no cover
+        return None
+
+
+_QK_QUANT_DTYPES = {
+    "int8": torch.int8,
+    "fp8_e4m3": torch.float8_e4m3fn,
+}
+
+
+@dataclass(frozen=True)
+class QuantConfig:
+    dtype_qk: str | None = None
+    q_block_size: int = 1
+    k_block_size: int = 16
+
+    @classmethod
+    def from_backend_kwargs(cls, backend_kwargs: dict | None) -> "QuantConfig":
+        q = (backend_kwargs or {}).get("quant") or {}
+        return cls(
+            dtype_qk=q.get("dtype_qk"),
+            q_block_size=int(q.get("q_block_size", 1) or 1),
+            k_block_size=int(q.get("k_block_size", 16) or 16),
+        )
+
+    @property
+    def enabled(self) -> bool:
+        return self.dtype_qk is not None
+
+    def quantize(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, quantize_fn):
+        qk_quant_dtype = _QK_QUANT_DTYPES[self.dtype_qk]
+        q_q, k_q, v_q, q_sfs, k_sfs, v_sfs = quantize_fn(
+            q,
+            k,
+            v,
+            q_block_size=self.q_block_size,
+            k_block_size=self.k_block_size,
+            qk_quant_dtype=qk_quant_dtype,
+        )
+        sage_attn_sfs = (q_sfs, k_sfs, None, v_sfs)
+        num_elts_per_sage_attn_blk = (self.q_block_size, self.k_block_size, 0, 1)
+        return q_q, k_q, v_q, sage_attn_sfs, num_elts_per_sage_attn_blk
+
+
 def _workspace_bytes() -> int:
     import vllm.envs as envs
 
@@ -129,6 +191,30 @@ class TrtllmAttentionImpl(AttentionImpl):
         self.skip = SkipSoftmaxConfig.from_backend_kwargs(backend_kwargs)
         self._warned_missing_timestep = False
 
+        self.quant = QuantConfig.from_backend_kwargs(backend_kwargs)
+        # Resolve the SAGE quantize fn once at init so the compiled forward path never calls the
+        # lru_cache-wrapped getter (which triggers a Dynamo graph break every step).
+        self._sage_quantize_fn = None
+        if self.quant.enabled:
+            if self.quant.dtype_qk not in _QK_QUANT_DTYPES:
+                raise RuntimeError(
+                    f"TRTLLM_ATTN quant (SAGE) supports dtype_qk in {sorted(_QK_QUANT_DTYPES)}, got "
+                    f"{self.quant.dtype_qk!r}. FLASHINFER_ATTN dtypes (float16/bfloat16) are not SAGE."
+                )
+            if not _sage_kernel_available():
+                raise RuntimeError(
+                    "TRTLLM_ATTN quant (SAGE) was requested but this FlashInfer build does not "
+                    "expose the trtllm-gen sage_attn_sfs kernel path. Install a FlashInfer build "
+                    "that provides it, or remove the quant config."
+                )
+            self._sage_quantize_fn = _sage_quantize_fn()
+            if self._sage_quantize_fn is None:
+                raise RuntimeError(
+                    "TRTLLM_ATTN quant (SAGE) was requested but this FlashInfer build lacks "
+                    "trtllm_sage_attention_quantize (added in flashinfer >= 0.6.16rc1). Upgrade "
+                    "FlashInfer, or remove the quant config."
+                )
+
     def set_layer_calibration(self, a: float, b: float) -> None:
         self.skip = replace(self.skip, a=a, b=b)
 
@@ -203,6 +289,15 @@ class TrtllmAttentionImpl(AttentionImpl):
 
         _skip_factor = self._resolve_skip_factor(kv_len)
 
+        # SAGE kwargs are only understood by newer FlashInfer builds; pass them exclusively when
+        # SAGE quant is active (which already requires the kernel, checked at init) so the dense
+        # path stays compatible with older builds that lack these parameters.
+        sage_kwargs: dict = {}
+        if self.quant.enabled:
+            q, k, v, sage_attn_sfs, sage_block_sizes = self.quant.quantize(q, k, v, self._sage_quantize_fn)
+            sage_kwargs["sage_attn_sfs"] = sage_attn_sfs
+            sage_kwargs["num_elts_per_sage_attn_blk"] = sage_block_sizes
+
         out = trtllm_ragged_attention_deepseek(
             query=q,
             key=k,
@@ -222,5 +317,6 @@ class TrtllmAttentionImpl(AttentionImpl):
             is_causal=self.causal,
             return_lse=False,
             skip_softmax_threshold_scale_factor=_skip_factor,
+            **sage_kwargs,
         )
         return out.reshape(batch, q_len, num_q_heads, head_dim)

vllm_omni/diffusion/data.py

@@ -1419,44 +1419,88 @@ class SkipSoftmaxSpec:
             raise ValueError("skip_softmax: set either target_sparsity or threshold, not both.")
 
 
+@dataclass
+class AttnQuantSpec:
+    dtype_qk: str | None = None
+    dtype_vo: str | None = None
+    q_block_size: int = 1
+    k_block_size: int = 16
+    flashinfer_backend: str | None = None
+
+    _VALID_DTYPES = frozenset({"float16", "bfloat16", "int8", "fp8_e4m3"})
+    _VALID_BLOCK_SIZES = frozenset({1, 4, 16})
+
+    def __post_init__(self) -> None:
+        for name, v in (("dtype_qk", self.dtype_qk), ("dtype_vo", self.dtype_vo)):
+            if v is not None and v not in self._VALID_DTYPES:
+                raise ValueError(f"quant.{name}={v!r} unsupported; use one of {sorted(self._VALID_DTYPES)}.")
+        for name, v in (("q_block_size", self.q_block_size), ("k_block_size", self.k_block_size)):
+            if v not in self._VALID_BLOCK_SIZES:
+                raise ValueError(
+                    f"quant.{name}={v!r} unsupported; kernels exist only for {sorted(self._VALID_BLOCK_SIZES)}."
+                )
+
+    @property
+    def enabled(self) -> bool:
+        return self.dtype_qk is not None or self.dtype_vo is not None
+
+
 @dataclass
 class AttentionSpec:
     """Specifies a backend and its typed backend-specific config for one attention role."""
 
     backend: str
     skip_softmax: SkipSoftmaxSpec | None = None
+    quant: AttnQuantSpec | None = None
     skip_calibration: dict | None = field(default=None, repr=False)
 
     def __post_init__(self) -> None:
         if not isinstance(self.backend, str):
             raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}")
-        self.skip_softmax = self._coerce_skip_softmax(self.skip_softmax)
+        self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax")
+        self.quant = self._coerce(self.quant, AttnQuantSpec, "quant")
         if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN":
             raise ValueError(
                 f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. "
-                f"Remove skip_softmax or set backend to TRTLLM_ATTN."
+                "Remove skip_softmax or set backend to TRTLLM_ATTN."
+            )
+        if self.quant is not None and self.backend.upper() not in ("TRTLLM_ATTN", "FLASHINFER_ATTN"):
+            raise ValueError(
+                f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but "
+                f"backend={self.backend!r}. Remove quant or set a supported backend."
             )
 
     @staticmethod
-    def _coerce_skip_softmax(value: Any) -> SkipSoftmaxSpec | None:
-        if value is None or isinstance(value, SkipSoftmaxSpec):
+    def _coerce(value: Any, cls: type, field_name: str) -> Any:
+        if value is None or isinstance(value, cls):
             return value
         if isinstance(value, Mapping):
-            return SkipSoftmaxSpec(**dict(value))
-        raise TypeError(f"Expected dict or SkipSoftmaxSpec for skip_softmax, got {type(value)!r}")
+            return cls(**dict(value))
+        raise TypeError(f"Expected dict or {cls.__name__} for {field_name}, got {type(value)!r}")
 
     def backend_kwargs(self) -> dict[str, Any] | None:
         """Serialize typed backend config into the kwargs dict the backend impl consumes."""
-        if self.skip_softmax is None:
-            return None
-        ss = self.skip_softmax
         kw: dict[str, Any] = {}
-        if ss.threshold is not None:
-            kw["skip_softmax_threshold"] = ss.threshold
-        if ss.target_sparsity is not None:
-            kw["target_sparsity"] = ss.target_sparsity
-        if ss.disabled_until_timestep:
-            kw["disabled_until_timestep"] = ss.disabled_until_timestep
+        if self.skip_softmax is not None:
+            ss = self.skip_softmax
+            if ss.threshold is not None:
+                kw["skip_softmax_threshold"] = ss.threshold
+            if ss.target_sparsity is not None:
+                kw["target_sparsity"] = ss.target_sparsity
+            if ss.disabled_until_timestep:
+                kw["disabled_until_timestep"] = ss.disabled_until_timestep
+        if self.quant is not None and self.quant.enabled:
+            q = self.quant
+            quant_kw: dict[str, Any] = {
+                "dtype_qk": q.dtype_qk,
+                "q_block_size": q.q_block_size,
+                "k_block_size": q.k_block_size,
+            }
+            if q.dtype_vo is not None:
+                quant_kw["dtype_vo"] = q.dtype_vo
+            if q.flashinfer_backend is not None:
+                quant_kw["flashinfer_backend"] = q.flashinfer_backend
+            kw["quant"] = quant_kw
         return kw or None
 
 
@@ -1526,7 +1570,7 @@ class AttentionConfig:
             normalized[role] = node
             return
 
-        spec_keys = {"backend", "skip_softmax"}
+        spec_keys = {"backend", "skip_softmax", "quant"}
         node_dict = dict(node)
         node_keys = set(node_dict)
         if node_keys & spec_keys: