From d12ea3e9ba9f8d76cd5fc1e9ce9b71681e2f3a2e Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Wed, 29 Jul 2026 03:02:28 -0700 Subject: [PATCH] docker: add Kimi K3 images (#32760) --- docker/kimi_k3/apply_deepep_k3_patch.sh | 169 + docker/kimi_k3/apply_deepgemm_situ_patch.py | 140 + .../kimi_k3/flashinfer-perkz-dcp-0.6.15.txt | 5639 +++++++++++++++++ docker/kimi_k3/kimi_k3_cu12.Dockerfile | 136 + docker/kimi_k3/kimi_k3_cu13.Dockerfile | 125 + docker/rocm.Dockerfile | 8 +- 6 files changed, 6213 insertions(+), 4 deletions(-) create mode 100755 docker/kimi_k3/apply_deepep_k3_patch.sh create mode 100644 docker/kimi_k3/apply_deepgemm_situ_patch.py create mode 100644 docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt create mode 100644 docker/kimi_k3/kimi_k3_cu12.Dockerfile create mode 100644 docker/kimi_k3/kimi_k3_cu13.Dockerfile diff --git a/docker/kimi_k3/apply_deepep_k3_patch.sh b/docker/kimi_k3/apply_deepep_k3_patch.sh new file mode 100755 index 000000000..6f4506196 --- /dev/null +++ b/docker/kimi_k3/apply_deepep_k3_patch.sh @@ -0,0 +1,169 @@ +#!/bin/bash +# Kimi-K3 DeepEP patch + rebuild -- IMAGE-BUILD variant of the runtime patcher. +# +# Stock DeepEP (deepseek-ai@d28bd67, shipped in the base image at +# /sgl-workspace/DeepEP) does not serve Kimi-K3: +# - LL topk capped at 11; K3 routes top-16 -> assert at internode_ll.cu +# - SWITCH_HIDDEN lacks 3584 (K3 latent-MoE dispatch dim) +# - normal-internode issues an unaligned 64-bit SourceMeta access for packed +# K3 FP8 scales when dispatch crosses the 8-rank NVL domain (EP > 8) +# - CUDA 13 relocated the cccl headers -> the stock build cannot find them +# +# Differences vs the devbox runtime script (why a separate copy): +# 1. arch comes from $TORCH_CUDA_ARCH_LIST (no GPU at image-build time; the +# runtime script probes torch.cuda.get_device_capability which needs a GPU) +# 2. adds the setup.py cccl include-dir fix (required to COMPILE on CUDA 13; +# the devbox had it applied manually) +# 3. adds the configs.cuh CPU/cycle timeout bump (cross-node init headroom) +# +# Idempotent (grep/count-guarded). The default wheel contains native cubins for +# Hopper sm_90, B200 sm_100a, and GB300 sm_103a. +set -euo pipefail +: "${TORCH_CUDA_ARCH_LIST:=9.0;10.0a;10.3a}" +DEEPEP_DIR="${DEEPEP_DIR:-/sgl-workspace/DeepEP}" +DEEPEP_COMMIT="${DEEPEP_COMMIT:-d28bd676c2120573c9f1425f0c16c39faa4117e6}" + +# PyTorch accepts CUDA architecture lists separated by spaces or semicolons. +read -r -a CUDA_ARCHES <<< "${TORCH_CUDA_ARCH_LIST//;/ }" +[ "${#CUDA_ARCHES[@]}" -gt 0 ] || { echo "ERROR: TORCH_CUDA_ARCH_LIST is empty"; exit 1; } +for CUDA_ARCH in "${CUDA_ARCHES[@]}"; do + [[ "$CUDA_ARCH" =~ ^[0-9]+\.[0-9]+a?$ ]] || { + echo "ERROR: unsupported CUDA architecture '$CUDA_ARCH' in TORCH_CUDA_ARCH_LIST='$TORCH_CUDA_ARCH_LIST'" + exit 1 + } +done + +# The base image is expected to ship the DeepEP source here; clone-pin if not. +if [ ! -d "$DEEPEP_DIR/csrc" ]; then + echo "== $DEEPEP_DIR missing — cloning deepseek-ai/DeepEP @ $DEEPEP_COMMIT" + git clone --recursive https://github.com/deepseek-ai/DeepEP.git "$DEEPEP_DIR" + git -C "$DEEPEP_DIR" checkout "$DEEPEP_COMMIT" + git -C "$DEEPEP_DIR" submodule update --init --recursive +fi +cd "$DEEPEP_DIR" + +echo "== [1/6] internode_ll.cu: LL topk caps 9/11 -> 16" +TOPK_CAPS_EXPECTED=$(grep -Eci "kNumMaxTop[Kk] = (9|11|16)" csrc/kernels/internode_ll.cu || true) +[ "$TOPK_CAPS_EXPECTED" -ge 2 ] || { + echo "ERROR: expected >=2 recognized topk caps, found $TOPK_CAPS_EXPECTED" + exit 1 +} +sed -i 's/constexpr int kNumMaxTopK = 11;/constexpr int kNumMaxTopK = 16;/' csrc/kernels/internode_ll.cu +sed -i 's/constexpr int kNumMaxTopK = 9;/constexpr int kNumMaxTopK = 16;/' csrc/kernels/internode_ll.cu +sed -i 's/constexpr int kNumMaxTopk = 9;/constexpr int kNumMaxTopk = 16;/' csrc/kernels/internode_ll.cu +sed -i 's/constexpr int kNumMaxTopk = 11;/constexpr int kNumMaxTopk = 16;/' csrc/kernels/internode_ll.cu +TOPK_CAPS_PATCHED=$(grep -ci "kNumMaxTop[Kk] = 16" csrc/kernels/internode_ll.cu || true) +TOPK_CAPS_UNPATCHED=$(grep -Eci "kNumMaxTop[Kk] = (9|11)" csrc/kernels/internode_ll.cu || true) +if [ "$TOPK_CAPS_PATCHED" -ne "$TOPK_CAPS_EXPECTED" ] || [ "$TOPK_CAPS_UNPATCHED" -ne 0 ]; then + echo "ERROR: expected $TOPK_CAPS_EXPECTED topk caps patched, found $TOPK_CAPS_PATCHED patched and $TOPK_CAPS_UNPATCHED unpatched" + exit 1 +fi + +echo "== [2/6] launch.cuh: SWITCH_HIDDEN += case 3584 (K3 latent MoE)" +grep -q "case_macro(3584)" csrc/kernels/launch.cuh || \ + sed -i 's/case 4096: case_macro(4096);/case 3584: case_macro(3584); case 4096: case_macro(4096);/' csrc/kernels/launch.cuh +grep -q "case_macro(3584)" csrc/kernels/launch.cuh || { echo "ERROR: hidden 3584 case not applied"; exit 1; } + +echo "== [3/6] configs.cuh: raise CPU/cycle timeouts 100s -> 1000s (cross-node init headroom)" +# '@' delimiter (pattern contains '#'); '$' anchor avoids the ENABLE_FAST_DEBUG '...10' line. +sed -i 's@^#define NUM_CPU_TIMEOUT_SECS 100$@#define NUM_CPU_TIMEOUT_SECS 1000@' csrc/kernels/configs.cuh +sed -i 's@^#define NUM_TIMEOUT_CYCLES 200000000000ull@#define NUM_TIMEOUT_CYCLES 2000000000000ull@' csrc/kernels/configs.cuh +grep -q "NUM_CPU_TIMEOUT_SECS 1000$" csrc/kernels/configs.cuh || echo " WARN: CPU timeout not raised (layout changed?) — non-fatal" + +echo "== [4/6] tests/test_low_latency.py: nvfp4 tolerance 0.007 -> 0.008 (topk16 noise)" +python3 - <<'EOF' +p = "tests/test_low_latency.py" +s = open(p).read() +old = "elif dispatch_use_nvfp4:\n diff_threshold = 0.007" +new = "elif dispatch_use_nvfp4:\n diff_threshold = 0.008" +if old in s: + open(p, "w").write(s.replace(old, new)); print(" patched") +else: + print(" already patched or layout changed (skipped)") +EOF + +echo "== [5/6] internode.cu: 4-byte-aligned SourceMeta scalar access (EP>8 normal internode)" +python3 - <<'EOF' +from pathlib import Path +path = Path("csrc/kernels/internode.cu") +source = path.read_text() +changed = False +abi_old = 'EP_STATIC_ASSERT(sizeof(SourceMeta) % sizeof(int) == 0, "Invalid size of `SourceMeta`");' +abi_new = 'EP_STATIC_ASSERT(sizeof(SourceMeta) == 2 * sizeof(int), "SourceMeta scalar access requires exactly two int fields");' +if source.count(abi_old) == 2 and source.count(abi_new) == 0: + source = source.replace(abi_old, abi_new); changed = True; print(" tightened both SourceMeta ABI assertions") +elif source.count(abi_old) == 0 and source.count(abi_new) == 2: + print(" ABI assertions already tightened") +else: + raise SystemExit(f"ERROR: unexpected SourceMeta ABI layout: old={source.count(abi_old)}, new={source.count(abi_new)}") +replacements = ( + ("sender store", + """ // Copy source metadata into symmetric send buffer + if (lane_id < num_topk_ranks) + st_na_global(reinterpret_cast(dst_send_buffers[lane_id]), src_meta); +""", + """ // SourceMeta may be only 4-byte aligned after packed scales. + // Store its two int fields separately to avoid an unaligned 64-bit store. + if (lane_id < num_topk_ranks) { + auto meta_values = reinterpret_cast(dst_send_buffers[lane_id]); + st_na_global(meta_values, src_meta.src_rdma_rank); + st_na_global(meta_values + 1, src_meta.is_token_in_nvl_rank_bits); + } +"""), + ("forwarder load", + """ auto src_meta = ld_nc_global(reinterpret_cast(shifted + hidden_bytes + scale_bytes)); +""", + """ auto src_meta_values = reinterpret_cast(shifted + hidden_bytes + scale_bytes); + SourceMeta src_meta; + src_meta.src_rdma_rank = ld_nc_global(src_meta_values); + src_meta.is_token_in_nvl_rank_bits = ld_nc_global(src_meta_values + 1); +"""), + ("receiver load", + """ auto meta = ld_nc_global(reinterpret_cast(shifted + hidden_bytes + scale_bytes)); +""", + """ auto meta_values = reinterpret_cast(shifted + hidden_bytes + scale_bytes); + SourceMeta meta; + meta.src_rdma_rank = ld_nc_global(meta_values); + meta.is_token_in_nvl_rank_bits = ld_nc_global(meta_values + 1); +"""), +) +for label, old, new in replacements: + oc, nc = source.count(old), source.count(new) + if oc == 1 and nc == 0: + source = source.replace(old, new); changed = True; print(f" patched {label}") + elif oc == 0 and nc == 1: + print(f" {label} already patched") + else: + raise SystemExit(f"ERROR: unexpected {label} layout: old={oc}, new={nc}") +for unsafe in ("st_na_global(reinterpret_cast", "ld_nc_global(reinterpret_cast"): + if unsafe in source: + raise SystemExit(f"ERROR: unsafe SourceMeta access remains: {unsafe}") +if changed: + path.write_text(source) +EOF + +echo "== [5b] setup.py: add CUDA 13 cccl include dir (compile fix)" +grep -q "/usr/local/cuda/include/cccl" setup.py || \ + sed -i "s#\( include_dirs = \['csrc/'\]\)#\1\n include_dirs.append('/usr/local/cuda/include/cccl')#" setup.py +grep -q "/usr/local/cuda/include/cccl" setup.py || { echo "ERROR: cccl include not added to setup.py"; exit 1; } + +echo "== [6/6] rebuild + reinstall for TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST (no GPU needed)" +rm -rf build dist +TORCH_CUDA_ARCH_LIST="$TORCH_CUDA_ARCH_LIST" python3 setup.py bdist_wheel +pip install dist/*.whl --force-reinstall --no-deps + +echo "== verify installed cubin arches" +SO=$(find /usr/local/lib/python3*/dist-packages -maxdepth 1 -name "deep_ep_cpp*.so" | head -1) +[ -n "$SO" ] || { echo "ERROR: installed deep_ep_cpp shared object not found"; exit 1; } +CUBIN_LIST=$(cuobjdump --list-elf "$SO") +VERIFIED_SMS=() +for CUDA_ARCH in "${CUDA_ARCHES[@]}"; do + SM="sm_${CUDA_ARCH/./}" + grep -Eq "(^|[^[:alnum:]_])${SM}([^[:alnum:]_]|$)" <<< "$CUBIN_LIST" || { + echo "ERROR: cubin arch $SM not found in $SO" + printf '%s\n' "$CUBIN_LIST" | head + exit 1 + } + VERIFIED_SMS+=("$SM") +done +echo "== OK: DeepEP rebuilt (topk16 + hidden3584 + SourceMeta align + cccl) for ${VERIFIED_SMS[*]}" diff --git a/docker/kimi_k3/apply_deepgemm_situ_patch.py b/docker/kimi_k3/apply_deepgemm_situ_patch.py new file mode 100644 index 000000000..3e978e427 --- /dev/null +++ b/docker/kimi_k3/apply_deepgemm_situ_patch.py @@ -0,0 +1,140 @@ +"""Patch deep_gemm's mega-MoE JIT header to support Kimi-K3 SiTU activation. + +Mechanism (no _C.so rebuild needed — the kernel body is a runtime-JIT header): +the host passes `activation='swiglu'` with the magic `activation_clamp = +0.03125` (2^-5: exactly representable, round-trips through the host's float +stringification, and no legitimate swiglu clamp uses it; the host asserts +clamp >= 0 so a negative sentinel is not possible). In-kernel: +kActivationClamp == 0.03125f selects SiTU with K3 constants baked in: + beta = 4.0, linear_beta = 25.0 (config activation_situ_{beta,linear_beta}) +SiTU(gate, up) = beta*tanh(gate/beta)*sigmoid(gate) * (linear_beta*tanh(up/linear_beta)) + +A distinct clamp value produces a distinct JIT template instantiation, so the +new variant compiles fresh; cached swiglu kernels are unaffected. If you edit +the SiTU math itself, clear /root/.cache/deep_gemm first (same sentinel value +would otherwise hit a stale cache entry). + +Idempotent; also migrates the deprecated negative-sentinel V1 patch. +Run on every node: python3 apply_deepgemm_situ_patch.py +""" + +P = "/usr/local/lib/python3.12/dist-packages/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh" + +OLD = """ // Apply SwiGLU: silu(gate) * up + // Gate/up pairs: (0, 2), (1, 3), (4, 6), (5, 7) + auto fp32_values = reinterpret_cast(values); + #pragma unroll + for (uint32_t k = 0; k < 2; ++ k) { + auto bf16_gate = __float22bfloat162_rn(make_float2(fp32_values[k * 4], fp32_values[k * 4 + 1])); + auto bf16_up = __float22bfloat162_rn(make_float2(fp32_values[k * 4 + 2], fp32_values[k * 4 + 3])); + + // Clamp + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) { + bf16_gate = __hmin2(bf16_gate, {kActivationClamp, kActivationClamp}); + bf16_up = __hmax2(bf16_up, {-kActivationClamp, -kActivationClamp}); + bf16_up = __hmin2(bf16_up, {kActivationClamp, kActivationClamp}); + } + + // SwiGLU + auto gate = __bfloat1622float2(bf16_gate); + auto neg_gate_exp = make_float2( + kFastMath ? __expf(-gate.x) : expf(-gate.x), + kFastMath ? __expf(-gate.y) : expf(-gate.y)); + const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp); + if constexpr (kFastMath) { + gate = __fmul2_rn(gate, {math::fast_rcp(denom.x), math::fast_rcp(denom.y)}); + } else { + gate = {gate.x / denom.x, gate.y / denom.y}; + } + const auto up = __bfloat1622float2(bf16_up); + activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights); + } +""" + +NEW = """ // Apply activation: SwiGLU, or Kimi-K3 SiTU via sentinel + // Gate/up pairs: (0, 2), (1, 3), (4, 6), (5, 7) + // K3-SITU-PATCH: kActivationClamp == 0.03125f (2^-5 magic; + // host asserts clamp >= 0 so negatives can't sentinel) selects SiTU: + // act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate) + // up' = kSituLinearBeta * tanh(up/kSituLinearBeta) + // K3 config constants baked in (activation_situ_{beta,linear_beta}). + constexpr bool kUseSitu = (kActivationClamp == 0.03125f); + constexpr float kSituBeta = 4.0f; + constexpr float kSituLinearBeta = 25.0f; + auto fp32_values = reinterpret_cast(values); + #pragma unroll + for (uint32_t k = 0; k < 2; ++ k) { + auto bf16_gate = __float22bfloat162_rn(make_float2(fp32_values[k * 4], fp32_values[k * 4 + 1])); + auto bf16_up = __float22bfloat162_rn(make_float2(fp32_values[k * 4 + 2], fp32_values[k * 4 + 3])); + + // Clamp (SwiGLU-with-limit only; SiTU soft-clips below) + if constexpr (!kUseSitu && kActivationClamp != cute::numeric_limits::infinity()) { + bf16_gate = __hmin2(bf16_gate, {kActivationClamp, kActivationClamp}); + bf16_up = __hmax2(bf16_up, {-kActivationClamp, -kActivationClamp}); + bf16_up = __hmin2(bf16_up, {kActivationClamp, kActivationClamp}); + } + + // sigmoid(gate) + auto gate = __bfloat1622float2(bf16_gate); + auto neg_gate_exp = make_float2( + kFastMath ? __expf(-gate.x) : expf(-gate.x), + kFastMath ? __expf(-gate.y) : expf(-gate.y)); + const auto denom = __fadd2_rn({1.0f, 1.0f}, neg_gate_exp); + float2 sig; + if constexpr (kFastMath) { + sig = {math::fast_rcp(denom.x), math::fast_rcp(denom.y)}; + } else { + sig = {1.0f / denom.x, 1.0f / denom.y}; + } + auto up = __bfloat1622float2(bf16_up); + if constexpr (kUseSitu) { + // K3-SITU-PATCH: tanh-bounded gate, soft-clipped up + gate = {kSituBeta * tanhf(gate.x / kSituBeta) * sig.x, + kSituBeta * tanhf(gate.y / kSituBeta) * sig.y}; + up = {kSituLinearBeta * tanhf(up.x / kSituLinearBeta), + kSituLinearBeta * tanhf(up.y / kSituLinearBeta)}; + } else { + // SwiGLU: silu(gate) * up + gate = __fmul2_rn(gate, sig); + } + activation_values[i][k] = __fmul2_rn(__fmul2_rn(gate, up), weights); + } +""" + +V1_LINES = """ constexpr bool kUseSitu = kActivationClamp < 0.0f; + constexpr float kSituBeta = kUseSitu ? -kActivationClamp : 1.0f; + constexpr float kSituLinearBeta = 25.0f;""" +V2_LINES = """ constexpr bool kUseSitu = (kActivationClamp == 0.03125f); + constexpr float kSituBeta = 4.0f; + constexpr float kSituLinearBeta = 25.0f;""" + +s = open(P).read() +if V2_LINES in s: + print("already patched (v2)") +elif V1_LINES in s: + # migrate deprecated negative-sentinel v1 -> magic-value v2 + assert s.count(V1_LINES) == 1 + s = s.replace(V1_LINES, V2_LINES) + s = s.replace( + """ // K3-SITU-PATCH: kActivationClamp < 0 selects SiTU: + // act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate) + // up' = kSituLinearBeta * tanh(up/kSituLinearBeta) + // with kSituBeta = -kActivationClamp (host passes -beta). +""", + """ // K3-SITU-PATCH: kActivationClamp == 0.03125f (2^-5 magic; + // host asserts clamp >= 0 so negatives can't sentinel) selects SiTU: + // act = kSituBeta * tanh(gate/kSituBeta) * sigmoid(gate) + // up' = kSituLinearBeta * tanh(up/kSituLinearBeta) + // K3 config constants baked in (activation_situ_{beta,linear_beta}). +""", + ) + open(P, "w").write(s) + print("migrated v1 -> v2") +elif OLD in s: + assert s.count(OLD) == 1 + open(P, "w").write(s.replace(OLD, NEW)) + print("patched") +else: + raise SystemExit( + "ERROR: expected SwiGLU epilogue block not found — header layout changed" + ) diff --git a/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt b/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt new file mode 100644 index 000000000..945fbe395 --- /dev/null +++ b/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt @@ -0,0 +1,5639 @@ +diff --git a/flashinfer/cute_dsl/attention/mla_dispatch.py b/flashinfer/cute_dsl/attention/mla_dispatch.py +index f41976df..c56b21c5 100644 +--- a/flashinfer/cute_dsl/attention/mla_dispatch.py ++++ b/flashinfer/cute_dsl/attention/mla_dispatch.py +@@ -9,8 +9,10 @@ controlled by the ``cute_dsl_impl`` kwarg, with three valid values: + + * ``"auto"`` (default) — library picks the right implementation. + Monolithic by default, automatically promoted to modular when the call +- uses a feature monolithic doesn't support (currently: ``sinks``). ++ uses a feature monolithic doesn't support (currently: ``sinks``). DCP ++ requires monolithic and therefore keeps the monolithic implementation. + * ``"modular"`` — strict. Always run the modular implementation. ++ Raises :class:`ValueError` when DCP is enabled. + * ``"monolithic"`` — strict. Always run the monolithic implementation; + raise :class:`ValueError` if the call uses any modular-only feature. + No silent fallback — the contract is "you asked for monolithic, you +@@ -45,6 +47,12 @@ _logged_impls: set[str] = set() + # Add new entries here as more variants are exposed through the standalone + # signature. + MODULAR_ONLY_KWARGS = ("sinks",) ++DCP_KWARGS = ( ++ "enable_dcp", ++ "cp_world", ++ "cp_rank", ++ "causal_seqlens_kv_global", ++) + + + def _has_modular_only_feature(kwargs: dict) -> Optional[str]: +@@ -56,6 +64,45 @@ def _has_modular_only_feature(kwargs: dict) -> Optional[str]: + return None + + ++def _validate_dcp_kwargs(kwargs: dict) -> bool: ++ """Validate dispatcher-level DCP selection and return whether it is enabled. ++ ++ Tensor shape, dtype, and device validation belongs to the monolithic ++ wrapper, which has access to the normalized query batch. The dispatcher ++ only enforces the static feature contract and prevents non-default DCP ++ arguments from being silently ignored. ++ """ ++ enable_dcp = kwargs.get("enable_dcp", False) ++ if not isinstance(enable_dcp, bool): ++ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") ++ ++ cp_world = kwargs.get("cp_world", 1) ++ cp_rank = kwargs.get("cp_rank", 0) ++ causal_seqlens_kv_global = kwargs.get("causal_seqlens_kv_global") ++ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: ++ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") ++ if not isinstance(cp_rank, int) or isinstance(cp_rank, bool): ++ raise TypeError(f"cp_rank must be an integer, got {type(cp_rank).__name__}") ++ if enable_dcp and not 0 <= cp_rank < cp_world: ++ raise ValueError( ++ f"cp_rank must satisfy 0 <= cp_rank < cp_world, got " ++ f"cp_rank={cp_rank}, cp_world={cp_world}" ++ ) ++ if not enable_dcp: ++ nondefault = [] ++ if cp_world != 1: ++ nondefault.append(f"cp_world={cp_world!r}") ++ if cp_rank != 0: ++ nondefault.append(f"cp_rank={cp_rank!r}") ++ if causal_seqlens_kv_global is not None: ++ nondefault.append("causal_seqlens_kv_global") ++ if nondefault: ++ raise ValueError( ++ "DCP arguments require enable_dcp=True; got " + ", ".join(nondefault) ++ ) ++ return enable_dcp ++ ++ + def _resolve_impl(*, requested: str, kwargs: dict) -> str: + """Map a user request and call kwargs to a concrete impl name. + +@@ -67,9 +114,19 @@ def _resolve_impl(*, requested: str, kwargs: dict) -> str: + f"Invalid cute_dsl_impl={requested!r}; expected one of {VALID_IMPLS}" + ) + ++ enable_dcp = _validate_dcp_kwargs(kwargs) + needs_modular = _has_modular_only_feature(kwargs) + ++ if enable_dcp and needs_modular is not None: ++ raise ValueError( ++ f"DCP cannot be combined with {needs_modular!r}: DCP requires the " ++ "monolithic CuTeDSL MLA implementation, while the requested feature " ++ "requires the modular implementation." ++ ) ++ + if requested == "auto": ++ if enable_dcp: ++ return "monolithic" + return "modular" if needs_modular is not None else _DEFAULT_IMPL + + if requested == "monolithic" and needs_modular is not None: +@@ -80,6 +137,13 @@ def _resolve_impl(*, requested: str, kwargs: dict) -> str: + f"right impl based on the call) or cute_dsl_impl='modular'." + ) + ++ if requested == "modular" and enable_dcp: ++ raise ValueError( ++ "cute_dsl_impl='modular' was requested with enable_dcp=True, but " ++ "DCP is only supported by the monolithic CuTeDSL MLA implementation. " ++ "Use cute_dsl_impl='auto' or cute_dsl_impl='monolithic'." ++ ) ++ + return requested # "modular" or "monolithic" + + +@@ -91,18 +155,19 @@ def cute_dsl_mla_decode(*args, cute_dsl_impl: str = "auto", **kwargs): + :func:`flashinfer.cute_dsl.attention.wrappers.batch_mla.cute_dsl_mla_decode` + (modular, supports ``sinks=``) and + :func:`flashinfer.cute_dsl.attention.monolithic.mla_decode.cute_dsl_mla_decode` +- (monolithic, no variant support) — their signatures are otherwise +- identical. ++ (monolithic, supports static DCP). Implementation-specific keyword ++ arguments are validated before dispatch and stripped only when their ++ default values make them irrelevant to the selected implementation. + + Parameters + ---------- + cute_dsl_impl : str, default ``"auto"`` + ``"auto"`` (default) lets the dispatcher pick: monolithic by + default, modular when the call uses a modular-only feature +- (currently ``sinks``). ``"modular"`` and ``"monolithic"`` are +- strict — the dispatcher will not silently switch implementations, +- and ``"monolithic"`` raises :class:`ValueError` if the call uses +- a modular-only feature. ++ (currently ``sinks``); DCP requires monolithic. ``"modular"`` and ++ ``"monolithic"`` are strict — the dispatcher will not silently switch ++ implementations, and raises :class:`ValueError` for an incompatible ++ feature request. + """ + impl = _resolve_impl(requested=cute_dsl_impl, kwargs=kwargs) + +@@ -122,5 +187,8 @@ def cute_dsl_mla_decode(*args, cute_dsl_impl: str = "auto", **kwargs): + kwargs = {k: v for k, v in kwargs.items() if k not in MODULAR_ONLY_KWARGS} + from .monolithic.mla_decode import cute_dsl_mla_decode as _impl + else: ++ # Default-valued DCP kwargs are accepted by the common dispatcher but ++ # are not part of the modular implementation's ABI. ++ kwargs = {k: v for k, v in kwargs.items() if k not in DCP_KWARGS} + from .wrappers.batch_mla import cute_dsl_mla_decode as _impl + return _impl(*args, **kwargs) +diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_decode.py b/flashinfer/cute_dsl/attention/monolithic/mla_decode.py +index cb491ae1..eb9be701 100644 +--- a/flashinfer/cute_dsl/attention/monolithic/mla_decode.py ++++ b/flashinfer/cute_dsl/attention/monolithic/mla_decode.py +@@ -32,6 +32,7 @@ from flashinfer.utils import device_support_pdl + + from .mla_decode_fp16 import BlackwellMultiHeadLatentAttentionForwardFP16 + from .mla_decode_fp8 import BlackwellMultiHeadLatentAttentionForwardFP8 ++from .mla_helpers import MAX_SPLITS, ceil_div, compute_q_tile_layout + from flashinfer.cute_dsl.utils import ( + _as_cute_dsl_workspace_i8, + get_max_active_clusters, +@@ -40,6 +41,60 @@ from flashinfer.cute_dsl.utils import ( + ) + + ++_CUDA_GRID_Y_MAX = 65_535 ++_REDUCER_D_TILE_CANDIDATES = (1, 2, 4) ++_STATIC_REDUCER_MAX_SPLITS = 32 ++ ++ ++def _validate_nonpersistent_grid_y( ++ batch_size: int, num_q_tiles: int, is_persistent: bool ++) -> None: ++ """Reject nonpersistent launch grids whose Y dimension CUDA cannot encode.""" ++ grid_y = batch_size * num_q_tiles ++ if not is_persistent and grid_y > _CUDA_GRID_Y_MAX: ++ raise ValueError( ++ "nonpersistent CuTeDSL MLA grid.y would be " ++ f"{grid_y}, exceeding the CUDA limit {_CUDA_GRID_Y_MAX} " ++ f"(batch_size={batch_size}, num_q_tiles={num_q_tiles})" ++ ) ++ ++ ++def _get_reducer_d_tiles( ++ batch_size: int, ++ seq_len_q: int, ++ num_heads: int, ++ num_sms: int, ++ effective_split_kv: int, ++) -> int: ++ """Choose output-side reducer parallelism for an underfilled row grid. ++ ++ A reducer CTA normally owns one D512 row. When the real row grid does not ++ cover one resident wave, compare the conservative wave count for one, two, ++ or four equal D bands and keep the smallest topology with the shortest ++ per-band critical path. Once rows already cover every SM, retain one CTA ++ per row and avoid duplicated LSE work. ++ """ ++ reducer_rows = batch_size * seq_len_q * num_heads ++ if ( ++ reducer_rows <= 0 ++ or num_sms <= 0 ++ or reducer_rows >= num_sms ++ or effective_split_kv <= 1 ++ ): ++ return 1 ++ ++ best_tiles = 1 ++ best_waves = ceil_div(reducer_rows, num_sms) ++ for d_tiles in _REDUCER_D_TILE_CANDIDATES[1:]: ++ if d_tiles > effective_split_kv: ++ continue ++ waves = ceil_div(reducer_rows * d_tiles, num_sms) ++ if waves * best_tiles < best_waves * d_tiles: ++ best_tiles = d_tiles ++ best_waves = waves ++ return best_tiles ++ ++ + @functools.cache + def _get_split_kv_and_workspace_size( + B: int, +@@ -47,22 +102,32 @@ def _get_split_kv_and_workspace_size( + H: int, + kv_lora_rank: int, + max_active_blocks: int, ++ max_seq_len: Optional[int] = None, + ) -> Tuple[int, int]: +- """Cache split_kv and workspace_size since they are deterministic for the same params.""" +- # When folding S_q into heads, the workspace dims are the effective dims +- # (num_heads * F, q_len // F). get_workspace_size already pads H<128 to +- # 128, so passing num_heads_eff and seq_len_q_eff yields the right size. ++ """Return the nonempty split count and its workspace requirement. ++ ++ The occupancy heuristic can request more splits than the kernel's uniform ++ contiguous partitioning actually uses. When ``max_seq_len`` is known, ++ normalize the candidate to the number of nonempty fixed-size K chunks so ++ the grid and reducer do not carry an empty final split. ++ """ ++ # Flatten (query_token, head) into one contiguous row space. A cooperative ++ # 2-CTA MMA tile owns 128 consecutive rows and may cross token boundaries; ++ # only the final query tile can be padded. + mma_qk_tile_m = 128 +- fold_sq_ratio = BlackwellMultiHeadLatentAttentionForwardFP16.compute_fold_sq_ratio( +- H, q_len, mma_qk_tile_m +- ) +- num_heads_eff = H * fold_sq_ratio +- seq_len_q_eff = q_len // fold_sq_ratio ++ mma_qk_tile_n = 128 ++ _, num_q_tiles, _ = compute_q_tile_layout(H, q_len, mma_qk_tile_m) + split_kv = BlackwellMultiHeadLatentAttentionForwardFP16.get_split_kv_simplified( +- B, seq_len_q_eff, max_active_blocks ++ B, num_q_tiles, max_active_blocks + ) ++ if max_seq_len is not None: ++ if max_seq_len <= 0: ++ raise ValueError(f"max_seq_len must be > 0, got {max_seq_len}") ++ k_tile_total = ceil_div(max_seq_len, mma_qk_tile_n) ++ k_tiles_per_split = ceil_div(k_tile_total, split_kv) ++ split_kv = ceil_div(k_tile_total, k_tiles_per_split) + workspace_size = BlackwellMultiHeadLatentAttentionForwardFP16.get_workspace_size( +- num_heads_eff, seq_len_q_eff, kv_lora_rank, B, split_kv, cutlass.Float32 ++ mma_qk_tile_m, num_q_tiles, kv_lora_rank, B, split_kv, cutlass.Float32 + ) + return split_kv, workspace_size + +@@ -79,8 +144,19 @@ def _check_can_implement( + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, ++ enable_dcp: bool = False, ++ cp_world: int = 1, + ) -> None: + """Check if the kernel supports the given configuration (cached).""" ++ if not isinstance(enable_dcp, bool): ++ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") ++ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: ++ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") ++ if not enable_dcp and cp_world != 1: ++ raise ValueError( ++ f"cp_world={cp_world} requires enable_dcp=True; disabled DCP uses cp_world=1" ++ ) ++ + mma_qk_tiler_mn = (128, 128) + mma_pv_tiler_mn = (128, 256) + +@@ -129,15 +205,20 @@ def _get_compiled_mla_kernel( + is_persistent: bool, + is_var_seq: bool, + is_var_split_kv: bool, ++ reducer_d_tiles: int = 1, ++ reducer_max_splits: int = MAX_SPLITS, + skip_correction_threshold: float = 0.0, + is_workspace_size_zero: bool = False, + enable_pdl: bool = False, ++ enable_dcp: bool = False, ++ cp_world: int = 1, + ) -> Callable: + """Compile and cache an MLA decode kernel. + + Returns a callable that accepts (q_latent, q_rope, c_latent, c_rope, + page_table, o, lse, workspace, split_kv_scalar, cache_seqs, +- block_split_kvs, softmax_scale_scalar, output_scale_scalar). ++ causal_seqlens_kv_global, cp_rank_scalar, block_split_kvs, ++ softmax_scale_scalar, output_scale_scalar). + + All scalar arguments must be pre-wrapped as Int32/Float32. + """ +@@ -157,14 +238,6 @@ def _get_compiled_mla_kernel( + cutlass_dtype = torch_to_cutlass_dtype(torch_dtype) + cutlass_out_dtype = torch_to_cutlass_dtype(torch_out_dtype) + +- # Derive the seq_len_q-into-heads fold factor. F > 1 means the kernel +- # repacks the [H, S_q] tile to [H*F, S_q/F] internally so MTP / spec-decoding +- # with H < 128 fully populates the 128-wide MMA-M tile. +- fold_sq_ratio = KernelClass.compute_fold_sq_ratio( +- num_heads, seq_len_q, mma_qk_tiler_mn[0] +- ) +- fold_sq = fold_sq_ratio > 1 +- + kernel_obj = KernelClass( + acc_dtype=cutlass.Float32, + lse_dtype=cutlass.Float32, +@@ -181,7 +254,10 @@ def _get_compiled_mla_kernel( + enable_pdl=enable_pdl, + num_heads=num_heads, + seq_len_q=seq_len_q, +- fold_sq=fold_sq, ++ reducer_d_tiles=reducer_d_tiles, ++ reducer_max_splits=reducer_max_splits, ++ enable_dcp=enable_dcp, ++ cp_world=cp_world, + ) + + # All dimensions as sym_int — this matches the original kernel's use of +@@ -271,6 +347,18 @@ def _get_compiled_mla_kernel( + (sym_batch,), + assumed_align=16, + ) ++ # DCP's causal boundary is global while cache_seqs remains the physical ++ # rank-local length used for paging and split-K traversal. Keep the tensor ++ # absent from the disabled specialization so no load or pointer argument ++ # reaches its generated device code. ++ if enable_dcp: ++ causal_seqlens_kv_global_fake = cute.runtime.make_fake_compact_tensor( ++ cutlass.Int32, ++ (sym_batch,), ++ assumed_align=4, ++ ) ++ else: ++ causal_seqlens_kv_global_fake = None + # block_split_kvs: [batch_size] — int32 (only needed for is_var_split_kv=True) + if is_var_split_kv: + block_split_kvs_fake = cute.runtime.make_fake_compact_tensor( +@@ -295,6 +383,8 @@ def _get_compiled_mla_kernel( + workspace_fake, + Int32(1), # split_kv placeholder + cache_seqs_fake, ++ causal_seqlens_kv_global_fake, ++ Int32(0), # cp_rank placeholder (runtime-uniform, not a JIT key) + block_split_kvs_fake, + Float32(1.0), # softmax_scale placeholder + Float32(1.0), # output_scale placeholder +@@ -323,6 +413,10 @@ def cute_dsl_mla_decode( + enable_pdl: Optional[bool] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, ++ enable_dcp: bool = False, ++ cp_world: int = 1, ++ cp_rank: int = 0, ++ causal_seqlens_kv_global: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + """CuTe DSL MLA decode kernel for Blackwell SM100. + +@@ -336,8 +430,9 @@ def cute_dsl_mla_decode( + Pre-allocated workspace buffer (int8 or uint8). Required size depends on batch size + and split_kv (auto-computed from B, q_len, and number of SMs): + +- - Formula: ``B * H * q_len * split_kv * (kv_lora_rank + 1) * 4`` bytes +- (0 when split_kv == 1, which happens when B >= num_SMs / 2) ++ - Formula: ``B * 128 * q_tiles * split_kv * (kv_lora_rank + 1) * 4`` ++ bytes, where ``q_tiles = ceil(q_len * H / 128)`` ++ (0 when ``split_kv == 1``, typically once ``B * q_tiles >= num_SMs / 2``) + - Typical max: ~18 MB on a 148-SM GPU (e.g. B=4..8, H=128, D=512) + - Safe default: 128 MB covers all realistic configurations + kv_lora_rank : int +@@ -355,7 +450,7 @@ def cute_dsl_mla_decode( + output_scale : float + Scale factor applied to the output. + out : Optional[torch.Tensor] +- Pre-allocated output tensor [B, q_len, H, kv_lora_rank]. ++ Contiguous pre-allocated output tensor [B, q_len, H, kv_lora_rank]. + out_dtype : Optional[torch.dtype] + Output data type. If None, defaults to torch.bfloat16 (matching trtllm-gen backend). + Supported values: torch.bfloat16, torch.float8_e4m3fn (FP8 input only), +@@ -375,13 +470,25 @@ def cute_dsl_mla_decode( + * ``[B * q_len, H]`` (matches ``trtllm-gen`` shape; the wrapper + reshapes via ``.view`` to the native layout). + ++ Caller-provided buffers must be contiguous. ++ + If ``return_lse`` is True and this is None, a buffer of the native + ``[B, q_len, H]`` shape is allocated internally. + return_lse : bool + Whether to return LSE values. When True, the function returns + ``(out, lse)`` (the ``lse`` tensor returned is in whatever shape + the caller supplied; if no ``lse`` was supplied, ``[B, q_len, H]``). +- ++ enable_dcp : bool ++ Enable static cyclic decode context-parallel masking. DCP returns a ++ rank-local attention state and therefore requires ``return_lse=True``. ++ cp_world : int ++ Compile-time context-parallel world size. The local cache owns global ++ token positions ``cp_world * local_k + cp_rank``. ++ cp_rank : int ++ Runtime-uniform context-parallel rank. ++ causal_seqlens_kv_global : Optional[torch.Tensor] ++ Contiguous CUDA int32 tensor ``[B]`` containing the global exclusive ++ causal bound for the newest query token. Required when DCP is enabled. + Returns + ------- + torch.Tensor or Tuple[torch.Tensor, torch.Tensor] +@@ -395,9 +502,88 @@ def cute_dsl_mla_decode( + assert kv_cache.dtype == query.dtype, ( + f"kv_cache dtype {kv_cache.dtype} must match query dtype {query.dtype}" + ) ++ if query.ndim != 4: ++ raise ValueError( ++ "query must have shape " ++ "[batch_size, q_len_per_request, num_heads, head_dim_qk]" ++ ) + B, q_len, H, D_qk = query.shape + assert D_qk == kv_lora_rank + qk_rope_head_dim + ++ if not isinstance(enable_dcp, bool): ++ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") ++ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: ++ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") ++ if not isinstance(cp_rank, int) or isinstance(cp_rank, bool): ++ raise TypeError(f"cp_rank must be an integer, got {type(cp_rank).__name__}") ++ ++ if enable_dcp: ++ if not 0 <= cp_rank < cp_world: ++ raise ValueError( ++ f"cp_rank must satisfy 0 <= cp_rank < cp_world, got " ++ f"cp_rank={cp_rank}, cp_world={cp_world}" ++ ) ++ if not return_lse: ++ raise ValueError( ++ "enable_dcp=True requires return_lse=True so rank-local states " ++ "can be merged with their LSE values" ++ ) ++ if causal_seqlens_kv_global is None: ++ raise ValueError( ++ "causal_seqlens_kv_global is required when enable_dcp=True" ++ ) ++ if not isinstance(causal_seqlens_kv_global, torch.Tensor): ++ raise TypeError( ++ "causal_seqlens_kv_global must be a torch.Tensor, got " ++ f"{type(causal_seqlens_kv_global).__name__}" ++ ) ++ if causal_seqlens_kv_global.dtype != torch.int32: ++ raise ValueError( ++ "causal_seqlens_kv_global must have dtype torch.int32, got " ++ f"{causal_seqlens_kv_global.dtype}" ++ ) ++ if not causal_seqlens_kv_global.is_cuda: ++ raise ValueError("causal_seqlens_kv_global must be a CUDA tensor") ++ if causal_seqlens_kv_global.device != query.device: ++ raise ValueError( ++ "causal_seqlens_kv_global must be on the query device " ++ f"{query.device}, got {causal_seqlens_kv_global.device}" ++ ) ++ if tuple(causal_seqlens_kv_global.shape) != (B,): ++ raise ValueError( ++ f"causal_seqlens_kv_global must have shape ({B},), got " ++ f"{tuple(causal_seqlens_kv_global.shape)}" ++ ) ++ if not causal_seqlens_kv_global.is_contiguous(): ++ raise ValueError("causal_seqlens_kv_global must be contiguous") ++ else: ++ nondefault = [] ++ if cp_world != 1: ++ nondefault.append(f"cp_world={cp_world}") ++ if cp_rank != 0: ++ nondefault.append(f"cp_rank={cp_rank}") ++ if causal_seqlens_kv_global is not None: ++ nondefault.append("causal_seqlens_kv_global") ++ if nondefault: ++ raise ValueError( ++ "DCP arguments require enable_dcp=True; got " + ", ".join(nondefault) ++ ) ++ ++ # The single Q TMA descriptor flattens (query_token, head) into one ++ # globally bounded affine row mode. Adjacent tokens must therefore follow ++ # all H head rows. Normal contiguous queries and last-dimension slices ++ # satisfy this; materialize uncommon token-strided/head-subset views. ++ if query.stride(-1) != 1 or query.stride(1) != H * query.stride(2): ++ query = query.contiguous() ++ ++ # O/LSE are compiled with compact layouts. In particular, a token-gapped ++ # view cannot be flattened across an M128 tile boundary and would be ++ # silently written through the gap, so reject noncompact buffers early. ++ if out is not None and not out.is_contiguous(): ++ raise ValueError(f"out must be contiguous, got strides {out.stride()}") ++ if lse is not None and not lse.is_contiguous(): ++ raise ValueError(f"lse must be contiguous, got strides {lse.stride()}") ++ + q_dtype = query.dtype + # Resolve output dtype: for FP8 input, default to bfloat16 (matching trtllm-gen backend); + # for FP16/BF16 input, default to same as input. Allow override via out_dtype or out tensor. +@@ -415,7 +601,8 @@ def cute_dsl_mla_decode( + kv_cache = kv_cache.squeeze(1) + page_size = kv_cache.shape[1] + +- # Split query into latent and rope components — keep contiguous [B, q_len, H, D]. ++ # Split query into latent and rope components. The slices share the flat ++ # affine token/head row stride established above. + # The kernel's __call__ reinterprets to [H, D, q_len, B] via zero-cost make_tensor. + q_latent_k = query[..., :kv_lora_rank] + q_rope_k = query[..., kv_lora_rank:] +@@ -434,9 +621,13 @@ def cute_dsl_mla_decode( + # Cached split_kv and workspace_size computation + max_active_blocks = get_num_sm(query.device) + split_kv, workspace_size = _get_split_kv_and_workspace_size( +- B, q_len, H, kv_lora_rank, max_active_blocks ++ B, q_len, H, kv_lora_rank, max_active_blocks, max_seq_len + ) + ++ _, num_q_tiles, _ = compute_q_tile_layout(H, q_len) ++ is_persistent = not is_var_seq ++ _validate_nonpersistent_grid_y(B, num_q_tiles, is_persistent) ++ + # Prepare workspace: the CUTE signature uses int8, while public FlashInfer + # workspace buffers are byte storage and may be uint8. + workspace_buffer = _as_cute_dsl_workspace_i8(workspace_buffer) +@@ -446,6 +637,14 @@ def cute_dsl_mla_decode( + f"need {workspace_size} bytes" + ) + is_workspace_size_zero = workspace_size == 0 ++ k_tile_total = ceil_div(max_seq_len, 128) ++ k_tiles_per_split = ceil_div(k_tile_total, split_kv) ++ effective_split_kv = ceil_div(k_tile_total, k_tiles_per_split) ++ reducer_d_tiles = ( ++ 1 ++ if is_workspace_size_zero ++ else _get_reducer_d_tiles(B, q_len, H, max_active_blocks, effective_split_kv) ++ ) + if is_workspace_size_zero: + workspace_bytes = None + else: +@@ -486,9 +685,6 @@ def cute_dsl_mla_decode( + block_split_kvs = None + skip_correction_threshold = 0.0 + +- # for fix-length, set is_persistent to True; otherwise, set to False. +- is_persistent = not is_var_seq +- + # Validate configuration (cached, negligible overhead after first call) + _check_can_implement( + torch_dtype=q_dtype, +@@ -501,6 +697,8 @@ def cute_dsl_mla_decode( + is_persistent=is_persistent, + is_var_seq=is_var_seq, + is_var_split_kv=is_var_split_kv, ++ enable_dcp=enable_dcp, ++ cp_world=cp_world, + ) + + enable_pdl = device_support_pdl(query.device) if enable_pdl is None else enable_pdl +@@ -519,9 +717,13 @@ def cute_dsl_mla_decode( + is_persistent=is_persistent, + is_var_seq=is_var_seq, + is_var_split_kv=is_var_split_kv, ++ reducer_d_tiles=reducer_d_tiles, ++ reducer_max_splits=_STATIC_REDUCER_MAX_SPLITS, + skip_correction_threshold=skip_correction_threshold, + is_workspace_size_zero=is_workspace_size_zero, + enable_pdl=enable_pdl, ++ enable_dcp=enable_dcp, ++ cp_world=cp_world, + ) + + # Call the kernel +@@ -536,6 +738,8 @@ def cute_dsl_mla_decode( + workspace_bytes, + Int32(split_kv), + cache_seqs, ++ causal_seqlens_kv_global, ++ Int32(cp_rank), + block_split_kvs, + Float32(softmax_scale), + Float32(output_scale), +diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py +index 8e33acee..8202fb3b 100644 +--- a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py ++++ b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp16.py +@@ -77,6 +77,7 @@ from cutlass.cutlass_dsl import BaseDSL + + from .mla_helpers import ( + ceil_div, ++ compute_q_tile_layout, + MAX_SPLITS, + LOG2_E, + MLAStaticTileScheduler, +@@ -108,15 +109,17 @@ launcher and ``flashinfer/cute_dsl/attention/mla_dispatch.py`` for impl selectio + + Constraints: + * Data type requirements: +- - Input/output: Float16 ++ - Input: Float16 or BFloat16 ++ - Output: Float16 or BFloat16 + - Accumulation and LSE: Float32 + * Fixed architecture parameters: +- - Number of attention heads: 128 ++ - Number of attention heads: 1-128 + - Latent dimension: 512 + - RoPE dimension: 64 + * Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) + * Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +-* Query sequence length must be 1-4 ++* Query sequence length must be positive; query (token, head) rows are packed ++ continuously into 128-row M tiles with a safely padded final tile + * Only supports 2-CTA instructions + * Variable sequence length requires page table storage enabled + """ +@@ -138,7 +141,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + enable_pdl: bool, + num_heads: int = 128, + seq_len_q: int = 1, +- fold_sq: bool = False, ++ reducer_d_tiles: int = 1, ++ reducer_max_splits: int = MAX_SPLITS, ++ enable_dcp: bool = False, ++ cp_world: int = 1, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + +@@ -164,17 +170,24 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + :type is_var_split_kv: bool + :param enable_pdl: Whether to use PDL + :type enable_pdl: bool +- :param num_heads: Number of attention heads (pre-fold). Used for the +- per-row spec-decoding (MTP) causal mask q_token_index computation. ++ :param num_heads: Number of attention heads. Defines the flattened ++ ``(query_token, head)`` row geometry and division-free causal mask. + :type num_heads: int +- :param seq_len_q: Query sequence length (pre-fold). Combined with +- ``num_heads`` to derive the per-row q_token used by the causal mask. ++ :param seq_len_q: Query sequence length. Combined with ``num_heads`` ++ to define the flattened query-row extent and causal boundary. + :type seq_len_q: int +- :param fold_sq: Whether to fold tokens of ``seq_len_q`` into the head +- dimension so the M tile becomes [F sub_q_tok][num_heads heads]. +- Required when ``num_heads < mma_qk_tiler_mn[0]`` and ``seq_len_q > 1`` +- so the M tile is fully populated. +- :type fold_sq: bool ++ :param reducer_d_tiles: Number of independent D bands reduced per row. ++ :type reducer_d_tiles: int ++ :param reducer_max_splits: Compile-time reducer capacity. Direct users ++ retain the generic 256-split capacity by default; callers choosing ++ a smaller specialization must cap runtime split-KV accordingly. ++ :type reducer_max_splits: int ++ :param enable_dcp: Whether to compile the decode-context-parallel ++ global-coordinate causal mask. ++ :type enable_dcp: bool ++ :param cp_world: Number of cyclic context-parallel shards. This is a ++ compile-time parameter when DCP is enabled. ++ :type cp_world: int + """ + + self.latent_dim = 512 +@@ -189,25 +202,37 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv ++ if not 1 <= reducer_max_splits <= MAX_SPLITS: ++ raise ValueError( ++ f"reducer_max_splits must be in [1, {MAX_SPLITS}], " ++ f"got {reducer_max_splits}" ++ ) ++ if is_var_split_kv and reducer_max_splits != MAX_SPLITS: ++ raise ValueError( ++ "variable split-KV requires the generic reducer capacity " ++ f"of {MAX_SPLITS}, got {reducer_max_splits}" ++ ) ++ self.reducer_max_splits = reducer_max_splits + self.enable_pdl = enable_pdl +- # Original (pre-fold) num_heads and seq_len_q used for per-row +- # spec-decoding (MTP) causal q_token_index computation. When fold_sq is +- # True the M tile is laid out as [F sub_q_tok][num_heads heads]; the +- # full q_tok for row r is blk_coord[1] * F + (r // num_heads). ++ if cp_world < 1: ++ raise ValueError(f"cp_world must be positive, got {cp_world}") ++ if not enable_dcp and cp_world != 1: ++ raise ValueError( ++ "cp_world must be 1 when DCP is disabled, " ++ f"got cp_world={cp_world}" ++ ) ++ self.enable_dcp = enable_dcp ++ self.cp_world = cp_world ++ # Flatten (query token, head) into the physical 128-wide MMA M mode. ++ # Token boundaries may cross a tile; only the final global tile can be ++ # partially populated and is padded by the Q TMA OOB fill. + self.num_heads = num_heads + self.seq_len_q = seq_len_q +- # fold_sq (caller-controlled): whether the folding code path is enabled. +- # fold_sq_ratio (derived): fold factor F ≥ 1; the largest divisor of +- # seq_len_q with num_heads * F ≤ M_tile and F ≤ seq_len_q. When the +- # caller passes fold_sq=False, the kernel ignores the ratio. +- # When fold_sq=True but the derived ratio is 1, the folding branch +- # is taken with F=1 (a no-op transform). +- self.fold_sq = fold_sq +- self.fold_sq_ratio = ( +- BlackwellMultiHeadLatentAttentionForwardFP16.compute_fold_sq_ratio( +- num_heads, seq_len_q, mma_qk_tiler_mn[0] +- ) +- ) ++ ( ++ self.total_q_rows, ++ self.num_q_tiles, ++ self.tail_q_rows, ++ ) = compute_q_tile_layout(num_heads, seq_len_q, mma_qk_tiler_mn[0]) + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), +@@ -215,6 +240,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 ++ if reducer_d_tiles not in (1, 2, 4): ++ raise ValueError(f"unsupported reducer_d_tiles={reducer_d_tiles}") ++ self.reducer_d_tiles = reducer_d_tiles ++ self.reducer_d_tile = self.latent_dim // reducer_d_tiles + mma_qk_tiler_k = self.rope_dim if self.seq_len_q == 1 else self.rope_dim * 2 + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], +@@ -311,6 +340,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], ++ causal_seqlens_kv_global: Optional[cute.Tensor], ++ cp_rank: cutlass.Int32, + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, +@@ -346,6 +377,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor ++ :param causal_seqlens_kv_global: Global exclusive causal KV bound for ++ the newest query on each batch item. Present only for DCP kernels. ++ :type causal_seqlens_kv_global: cute.Tensor ++ :param cp_rank: Runtime context-parallel rank. ++ :type cp_rank: cutlass.Int32 + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax +@@ -386,6 +422,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + q_latent = _reinterpret_4d(q_latent) + q_rope = _reinterpret_4d(q_rope) + o = _reinterpret_4d(o) ++ # Keep unpacked output aliases for the split-KV reducer, which ++ # launches only real (head, q_token) rows. ++ o_unpacked = o + + # Reinterpret contiguous [num_pages, page_size, D] as [page_size, D, num_pages] + # Input stride: (PS*D, D, 1) → Target: (D, 1, PS*D) +@@ -419,54 +458,62 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + stride=(lse.stride[2], lse.stride[1], lse.stride[0]), + ), + ) ++ lse_unpacked = lse + +- # When num_heads < M tile, fold up to F = fold_sq_ratio tokens of +- # seq_len_q into the head dimension so M_eff = num_heads * F (≤ M_tile). +- # E.g., H=32, S_q=4 → F=4, M_eff=128, S_q_eff=1 +- # E.g., H=32, S_q=8 → F=4, M_eff=128, S_q_eff=2 +- # This works because MLA shares KV across all heads/queries independently. +- # Tensor layout: [H, D, S_q, B] → [H*F, D, S_q/F, B]; relies on +- # stride_S == stride_H * H (always true for contiguous [B, S_q, H, D] +- # tensors after _reinterpret_4d). +- if cutlass.const_expr(self.fold_sq): +- F = self.fold_sq_ratio +- +- def _fold_sq_4d(t): +- return cute.make_tensor( +- t.iterator, +- cute.make_layout( +- ( +- t.shape[0] * F, +- t.shape[1], +- t.shape[2] // F, +- t.shape[3], +- ), +- stride=( +- t.stride[0], +- t.stride[1], +- t.stride[2] * F, +- t.stride[3], +- ), +- ), +- ) +- +- q_latent = _fold_sq_4d(q_latent) +- q_rope = _fold_sq_4d(q_rope) +- o = _fold_sq_4d(o) +- # LSE: [H, S_q, B] → [H*F, S_q/F, B] +- lse = cute.make_tensor( +- lse.iterator, ++ # Flatten token-major [H, D, S_q, B] query storage into one globally ++ # bounded [S_q * H, D, B] M mode. Consecutive M128 tiles can cross ++ # token boundaries naturally, while TMA zero-fills only the final tail. ++ def _flatten_q_rows(t): ++ return cute.make_tensor( ++ t.iterator, + cute.make_layout( +- (lse.shape[0] * F, lse.shape[1] // F, lse.shape[2]), +- stride=(lse.stride[0], lse.stride[1] * F, lse.stride[2]), ++ (self.total_q_rows, t.shape[1], t.shape[3]), ++ stride=(t.stride[0], t.stride[1], t.stride[3]), + ), + ) + ++ q_latent = _flatten_q_rows(q_latent) ++ q_rope = _flatten_q_rows(q_rope) ++ ++ # O/LSE expose physical M128 query tiles to the main epilogue. Their ++ # final virtual padding is never dereferenced because stores are ++ # predicated by get_valid_q_rows(). ++ m_tile = self.mma_qk_tiler_mn[0] ++ runtime_num_q_tiles = cute.ceil_div(o.shape[0] * o.shape[2], m_tile) ++ o = cute.make_tensor( ++ o.iterator, ++ cute.make_layout( ++ ( ++ m_tile, ++ o.shape[1], ++ runtime_num_q_tiles, ++ o.shape[3], ++ ), ++ stride=( ++ o.stride[0], ++ o.stride[1], ++ o.stride[0] * m_tile, ++ o.stride[3], ++ ), ++ ), ++ ) ++ lse = cute.make_tensor( ++ lse.iterator, ++ cute.make_layout( ++ (m_tile, runtime_num_q_tiles, lse.shape[2]), ++ stride=( ++ lse.stride[0], ++ lse.stride[0] * m_tile, ++ lse.stride[2], ++ ), ++ ), ++ ) ++ + acc_o, acc_lse = self.initialize_workspace( +- q_latent.shape[0], +- q_latent.shape[1], +- q_latent.shape[2], +- q_latent.shape[3], ++ o.shape[0], ++ o.shape[1], ++ o.shape[2], ++ o.shape[3], + split_kv, + self.acc_dtype, + workspace, +@@ -732,6 +779,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + acc_lse, + split_kv, + cache_seqs, ++ causal_seqlens_kv_global, ++ cp_rank, + block_split_kvs, + softmax_scale_log2, + output_scale, +@@ -756,17 +805,21 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( +- o, +- lse, ++ o_unpacked, ++ lse_unpacked, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( +- grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), ++ grid=( ++ o_unpacked.shape[0] * self.reducer_d_tiles, ++ o_unpacked.shape[2], ++ o_unpacked.shape[3], ++ ), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], +- smem=MAX_SPLITS * self.acc_dtype.width // 8, ++ smem=self.reducer_max_splits * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + use_pdl=self.enable_pdl, +@@ -832,6 +885,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, ++ causal_seqlens_kv_global: Optional[cute.Tensor], ++ cp_rank: cutlass.Int32, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, +@@ -872,8 +927,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom +- :param mKR: Compressed rope tensor +- :type mKR: cute.Tensor ++ :param mQR: query rope tensor ++ :type mQR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor +@@ -896,6 +951,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor ++ :param causal_seqlens_kv_global: Per-batch global exclusive causal KV ++ bound for the newest query, present only when DCP is enabled. ++ :type causal_seqlens_kv_global: cute.Tensor ++ :param cp_rank: Runtime context-parallel rank. ++ :type cp_rank: cutlass.Int32 + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax +@@ -1273,6 +1333,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: ++ causal_global = cutlass.Int32(0) ++ if cutlass.const_expr(self.enable_dcp): ++ causal_global = causal_seqlens_kv_global[blk_coord[2]] + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, +@@ -1281,6 +1344,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], ++ causal_global=causal_global, ++ cp_rank=cp_rank, + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, +@@ -1337,6 +1402,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: ++ causal_global = cutlass.Int32(0) ++ if cutlass.const_expr(self.enable_dcp): ++ causal_global = causal_seqlens_kv_global[blk_coord[2]] + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, +@@ -1345,8 +1413,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], ++ causal_global=causal_global, ++ cp_rank=cp_rank, ++ split_start_key=k_index * self.mma_qk_tiler[1], + L=mCL.shape[1], +- H=mQL.shape[0], ++ H=self.get_valid_q_rows(blk_coord[1]), + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, +@@ -1366,11 +1437,74 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) ++ elif cutlass.const_expr(self.enable_dcp): ++ self.store_empty_dcp_work( ++ mO, ++ mLSE, ++ mAccO, ++ mAccLSE, ++ blk_coord, ++ self.get_valid_q_rows(blk_coord[1]), ++ tidx, ++ ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + return + ++ @cute.jit ++ def store_empty_dcp_work( ++ self, ++ mO: Optional[cute.Tensor], ++ mLSE: Optional[cute.Tensor], ++ mAccO: Optional[cute.Tensor], ++ mAccLSE: Optional[cute.Tensor], ++ blk_coord: cute.Coord, ++ valid_q_rows: cutlass.Int32, ++ tidx: cutlass.Int32, ++ ): ++ """Materialize the neutral attention state for a split with no K tile.""" ++ cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] ++ compute_threads = self.num_compute_warps * self.threads_per_warp ++ local_tidx = tidx % compute_threads ++ cta_row_base = blk_coord[0] * cta_m_rows ++ ++ for linear_idx in cutlass.range( ++ local_tidx, cta_m_rows * self.latent_dim, compute_threads ++ ): ++ cta_row = linear_idx // self.latent_dim ++ d_idx = linear_idx % self.latent_dim ++ q_row = cta_row_base + cta_row ++ if cute.elem_less(q_row, valid_q_rows): ++ if cutlass.const_expr(mAccO is None): ++ mO[q_row, d_idx, blk_coord[1], blk_coord[2]] = self.o_dtype(0.0) ++ else: ++ mAccO[ ++ q_row, ++ blk_coord[3], ++ d_idx, ++ blk_coord[1], ++ blk_coord[2], ++ ] = self.acc_dtype(0.0) ++ ++ if cute.elem_less(local_tidx, cta_m_rows): ++ q_row = cta_row_base + local_tidx ++ if cute.elem_less(q_row, valid_q_rows): ++ if cutlass.const_expr(mAccLSE is None): ++ mLSE[q_row, blk_coord[1], blk_coord[2]] = -self.lse_dtype.inf ++ else: ++ mAccLSE[ ++ q_row, blk_coord[3], blk_coord[1], blk_coord[2] ++ ] = -self.lse_dtype.inf ++ ++ @cute.jit ++ def get_valid_q_rows(self, q_tile_idx: cutlass.Int32) -> cutlass.Int32: ++ """Number of valid flattened query rows in one physical M tile.""" ++ valid_rows = self.mma_qk_tiler_mn[0] ++ if q_tile_idx == self.num_q_tiles - 1: ++ valid_rows = self.tail_q_rows ++ return valid_rows ++ + @cute.kernel + def reduction_kernel( + self, +@@ -1402,27 +1536,47 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() +- blk_coord = (bidx, bidy, bidz) ++ # Reducer blocks cover one D band of one real output row. Splitting ++ # D512 into independent D128 bands exposes enough CTAs for small-batch ++ # decode without any cross-CTA reduction or synchronization. ++ d_tile_idx = bidx % self.reducer_d_tiles ++ blk_coord = (bidx // self.reducer_d_tiles, bidy, bidz) ++ flat_q_row = blk_coord[1] * self.num_heads + blk_coord[0] ++ # The physical M tile is fixed at 128 rows, so map with shift/mask. ++ q_tile = flat_q_row >> 7 ++ q_tile_row = flat_q_row & 127 + local_split_kv = ( + block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv + ) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) +- k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) +- local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) ++ if cutlass.const_expr(self.enable_dcp): ++ k_tile_per_cta = cutlass.max( ++ cute.ceil_div(k_tile_total, local_split_kv), cutlass.Int32(1) ++ ) ++ local_split_kv = cutlass.max( ++ cute.ceil_div(k_tile_total, k_tile_per_cta), cutlass.Int32(1) ++ ) ++ else: ++ k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) ++ local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() +- storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) ++ storage = smem.allocate(self.reducer_max_splits * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) +- smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) ++ smem_lse_scale = cute.make_tensor( ++ lse_scale_ptr, cute.make_layout(self.reducer_max_splits) ++ ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() +- gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] ++ gLSE = mAccLSE[q_tile_row, None, q_tile, blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) +- lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) ++ lse_per_thread = cute.ceil_div( ++ self.reducer_max_splits, self.threads_per_warp ++ ) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype +@@ -1439,39 +1593,61 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) +- lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 ++ if cutlass.const_expr(self.enable_dcp): ++ has_valid_lse = lse_max != -self.lse_dtype.inf ++ lse_max = lse_max if has_valid_lse else 0.0 ++ else: ++ lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse +- global_lse = ( +- lse_max + cute.math.log2(sum_lse, fastmath=True) +- if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse # noqa: SIM201 +- else self.lse_dtype.inf +- ) +- if tidx == 0: +- # Convert from kernel-internal log2 base to the natural-log +- # convention exposed to callers (matches trtllm-gen / flash-attn). +- # `1.0 / LOG2_E == ln(2)`. +- mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( +- 1.0 / LOG2_E ++ if cutlass.const_expr(self.enable_dcp): ++ global_lse = ( ++ lse_max + cute.math.log2(sum_lse, fastmath=True) ++ if has_valid_lse ++ else -self.lse_dtype.inf + ) ++ else: ++ global_lse = ( ++ lse_max + cute.math.log2(sum_lse, fastmath=True) ++ if not sum_lse == self.lse_dtype(0.0) # noqa: SIM201 ++ or sum_lse != sum_lse ++ else self.lse_dtype.inf ++ ) ++ if d_tile_idx == 0: ++ if tidx == 0: ++ # Convert from kernel-internal log2 base to the natural-log ++ # convention exposed to callers (matches trtllm-gen / flash-attn). ++ # `1.0 / LOG2_E == ln(2)`. ++ mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( ++ 1.0 / LOG2_E ++ ) + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): +- smem_lse_scale[split_kv_idx] = cute.math.exp2( +- local_lse[i] - global_lse, fastmath=True +- ) ++ if cutlass.const_expr(self.enable_dcp): ++ smem_lse_scale[split_kv_idx] = ( ++ cute.math.exp2( ++ local_lse[i] - global_lse, fastmath=True ++ ) ++ if has_valid_lse ++ else self.acc_dtype(0.0) ++ ) ++ else: ++ smem_lse_scale[split_kv_idx] = cute.math.exp2( ++ local_lse[i] - global_lse, fastmath=True ++ ) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( +- self.latent_dim, self.threads_per_warp * self.num_compute_warps ++ self.reducer_d_tile, self.threads_per_warp * self.num_compute_warps + ) +- gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] ++ gAccO = mAccO[q_tile_row, None, None, q_tile, blk_coord[2]] + rAccO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.acc_dtype + ) +@@ -1479,11 +1655,19 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): +- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps ++ element_idx = ( ++ d_tile_idx * self.reducer_d_tile ++ + tidx ++ + j * self.threads_per_warp * self.num_compute_warps ++ ) + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): +- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps ++ element_idx = ( ++ d_tile_idx * self.reducer_d_tile ++ + tidx ++ + j * self.threads_per_warp * self.num_compute_warps ++ ) + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() +@@ -1665,7 +1849,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + +- # Flatten divide and partition global tensors for QK TMA load ++ # The single flattened descriptor is bounded by S_q * H. It can load ++ # across token boundaries and zero-fills only the final M128 tail. + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) +@@ -1677,7 +1862,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + ) + tSgQL = thr_mma_qk.partition_A(gQL) + tSgQR = thr_mma_qk.partition_A(gQR) +- + cta_m = min( + qk_params.tiled_mma_qk.op.shape_mnk[0] + // qk_params.tiled_mma_qk.thr_id.shape, +@@ -1725,7 +1909,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + cute.group_modes(qk_params.sQ_rope, 0, 3), + cute.group_modes(tSgQR, 0, 3), + ) +- + tKCsKC, tCLgCL = cpasync.tma_partition( + qk_params.tma_atom_c_latent, + 0, +@@ -1742,11 +1925,14 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + tSgKR, + ) + ++ # Q is flattened as (Sq * Hq, D, B). The logical query-tile ++ # coordinate therefore lives in the residual-M mode created by ++ # flat_divide, rather than in a separate grouped-query mode. + tQLgQL = tQLgQL_mkl[ +- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] ++ None, common_params.blk_coord[1], None, common_params.blk_coord[2] + ] + tQRgQR = tQRgQR_mkl[ +- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] ++ None, common_params.blk_coord[1], None, common_params.blk_coord[2] + ] + + # Flatten divide and partition global tensors for V TMA load +@@ -1837,6 +2023,33 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + load_pt_release_state, + ) + ++ @cute.jit ++ def issue_q_tma_load( ++ self, ++ tma_atom_q_latent: cute.CopyAtom, ++ tQLgQL: cute.Tensor, ++ tQsQ: cute.Tensor, ++ tma_atom_q_rope: cute.CopyAtom, ++ tQRgQR: cute.Tensor, ++ tQsQ_rope: cute.Tensor, ++ tma_bar_ptr, ++ ): ++ """Issue one bounded flattened query-tile load into shared Q stages.""" ++ for i in cutlass.range(self.iterations_qk_latent): ++ cute.copy( ++ tma_atom_q_latent, ++ tQLgQL[None, i], ++ tQsQ[None, (i, 0)], ++ tma_bar_ptr=tma_bar_ptr, ++ ) ++ for i in cutlass.range(self.iterations_qk_rope): ++ cute.copy( ++ tma_atom_q_rope, ++ tQRgQR[None, i], ++ tQsQ_rope[None, i], ++ tma_bar_ptr=tma_bar_ptr, ++ ) ++ + @cute.jit + def load_tma_qk_one_k_tile( + self, +@@ -1893,22 +2106,15 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + tma_bar_ptr = common_params.load_q_pipeline.producer_get_barrier( + load_q_producer_state + ) +- for i in cutlass.range(self.iterations_qk_latent): +- # load q latent +- cute.copy( +- qk_params.tma_atom_q_latent, +- qk_params.tQLgQL[None, 0, i], +- qk_params.tQsQ[None, (i, 0)], +- tma_bar_ptr=tma_bar_ptr, +- ) +- for i in cutlass.range(self.iterations_qk_rope): +- # load q rope +- cute.copy( +- qk_params.tma_atom_q_rope, +- qk_params.tQRgQR[None, 0, i], +- qk_params.tQsQ_rope[None, i], +- tma_bar_ptr=tma_bar_ptr, +- ) ++ self.issue_q_tma_load( ++ qk_params.tma_atom_q_latent, ++ qk_params.tQLgQL, ++ qk_params.tQsQ, ++ qk_params.tma_atom_q_rope, ++ qk_params.tQRgQR, ++ qk_params.tQsQ_rope, ++ tma_bar_ptr, ++ ) + load_q_producer_state.advance() + load_kv_pipeline = common_params.load_kv_pipeline + tma_bar_ptr = load_kv_pipeline.producer_get_barrier(load_kv_producer_state) +@@ -2361,24 +2567,48 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState] + """ + +- k_tile_total = cute.ceil_div(common_params.K, self.mma_qk_tiler[1]) +- + row_max = -self.acc_dtype.inf + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) + common_params.p_cor_pipeline.producer_acquire(p_cor_producer_state) + +- # Number of tiles from the global-K end that may contain causal-masked +- # positions. Min k_bound = K - (S_q-1), which can span up to +- # ceil((seq_len_q-2)/tile_N)+1 tiles (tile-boundary-crossing case). For +- # S_q=1 this reduces to 1 tile — identical to a plain K-bound check. ++ # The first tile that can contain a key at or beyond this query tile's ++ # earliest causal bound. A flat M tile starts at row q_tile * M, so ++ # floor(row / H) is its earliest query token. Later query tiles can ++ # therefore keep more fully dense K tiles on the compile-time unmasked ++ # path. The token division is outside the per-score loop (and is ++ # replicated by each participating CTA/compute group); tile_n is a ++ # static power of two, so the remaining division is a shift. + tile_n = self.mma_qk_tiler[1] +- mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 +- +- # first_mask_tile_idx is the global index of the first tile that may +- # need masking. Runtime because it depends on K (per-batch in +- # var-seq / split-KV). +- first_mask_tile_idx = k_tile_total - mask_tile_count ++ first_q_token = cutlass.Int32(0) ++ if cutlass.const_expr(self.num_q_tiles > 1): ++ first_q_token = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] ++ ) // self.num_heads ++ if cutlass.const_expr(self.enable_dcp): ++ # The earliest token in this full M128 query tile determines the ++ # dense-prefix boundary shared by both CTAs. Keep the physical ++ # local K bound separate so a partial final local tile is never ++ # misclassified as dense. ++ earliest_bound_numer = ( ++ common_params.causal_global ++ - common_params.cp_rank ++ - (self.seq_len_q - 1) ++ + first_q_token ++ ) ++ earliest_local_bound = cute.ceil_div( ++ cutlass.max(earliest_bound_numer, cutlass.Int32(0)), ++ self.cp_world, ++ ) ++ effective_local_bound = cutlass.min( ++ common_params.K, earliest_local_bound ++ ) ++ first_mask_tile_idx = effective_local_bound // tile_n ++ else: ++ first_mask_tile_idx = cutlass.max( ++ (common_params.K - self.seq_len_q + 1 + first_q_token) // tile_n, ++ cutlass.Int32(0), ++ ) + + # Phase 1: pure unmasked bulk tiles (all columns strictly < min k_bound). + while k_tile_count > 1 and k_index < first_mask_tile_idx: +@@ -2431,52 +2661,28 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + k_index = k_index + 1 + k_tile_count = k_tile_count - 1 + +- # Phase 3: this work-split's final tile. +- if cutlass.const_expr(common_params.mAccO is not None): +- # Split-KV: only apply mask when this final tile is globally in +- # the mask region (covers both last-split last-tile and straddling +- # splits). Runtime comparison. +- ( +- mma_s_consumer_state, +- p_mma_producer_state, +- p_cor_producer_state, +- row_max, +- row_sum, +- correction_factor, +- ) = self.softmax( +- common_params, +- softmax_params, +- k_index, +- mma_s_consumer_state, +- p_mma_producer_state, +- p_cor_producer_state, +- row_max, +- row_sum, +- correction_factor, +- k_index >= first_mask_tile_idx, +- True, +- ) +- else: +- ( +- mma_s_consumer_state, +- p_mma_producer_state, +- p_cor_producer_state, +- row_max, +- row_sum, +- correction_factor, +- ) = self.softmax( +- common_params, +- softmax_params, +- k_index, +- mma_s_consumer_state, +- p_mma_producer_state, +- p_cor_producer_state, +- row_max, +- row_sum, +- correction_factor, +- True, +- True, +- ) ++ # Phase 3: this work-split's final tile. Mask it only when its global ++ # index reaches the causal/K-bound region. ++ ( ++ mma_s_consumer_state, ++ p_mma_producer_state, ++ p_cor_producer_state, ++ row_max, ++ row_sum, ++ correction_factor, ++ ) = self.softmax( ++ common_params, ++ softmax_params, ++ k_index, ++ mma_s_consumer_state, ++ p_mma_producer_state, ++ p_cor_producer_state, ++ row_max, ++ row_sum, ++ correction_factor, ++ k_index >= first_mask_tile_idx, ++ True, ++ ) + + return mma_s_consumer_state, p_mma_producer_state, p_cor_producer_state + +@@ -2635,8 +2841,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param apply_mask: Whether the tile needs K-bound / causal masking (Python bool +- for the unmasked/masked bulk loops; runtime cutlass.Boolean for the +- split-KV final iter where mask only applies on the global last tile). ++ for the unmasked/masked bulk loops; runtime cutlass.Boolean for a ++ work-split final tile at the exact per-query-tile mask boundary). + :type apply_mask: bool | cutlass.Boolean + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean +@@ -2680,40 +2886,49 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max +- # Spec-decoding (MTP) causal mask: each row represents one (q_token, head) +- # pair; row r's effective K bound is K - (S_q - 1 - q_tok(r)). +- # With fold factor F = self.fold_sq_ratio (fold_sq=True), the M tile is +- # laid out as [F sub_q_tok][num_heads heads] and there are S_q/F outer +- # chunks indexed by blk_coord[1]: +- # q_tok(r) = blk_coord[1] * F + (r_global // num_heads) +- # r_global = row_in_cta + cluster_idx * (M_tile / cluster_m) +- # When fold_sq=False this reduces to q_tok = blk_coord[1]. For S_q=1 +- # this further reduces to k_bound = K (plain K-bound check). +- # Masked positions are filled with a large negative sentinel (not -inf) +- # to avoid NaN propagation when an entire row becomes masked. ++ # Spec-decoding (MTP) causal mask. A flattened row r corresponds to ++ # q_token=floor(r/H), whose last valid key is K-S_q+q_token. Avoid the ++ # per-element integer division using the equivalent integer predicate: ++ # ++ # H * (key_pos - K + S_q) <= r ++ # ++ # This arithmetic remains inside `if apply_mask`, so the compile-time ++ # dense path emits no row-mask work. Masked positions use a large ++ # negative sentinel (not -inf) to avoid all-masked-row NaNs. + cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if apply_mask: +- if cutlass.const_expr(self.fold_sq): +- q_tok = ( +- common_params.blk_coord[1] * self.fold_sq_ratio +- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) +- // self.num_heads ++ flat_q_row = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] ++ + common_params.blk_coord[0] * cta_m_rows ++ + tTR_tS[i][0] ++ ) ++ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index ++ if cutlass.const_expr(self.enable_dcp): ++ mask_threshold = self.num_heads * ( ++ self.cp_world * key_pos ++ + common_params.cp_rank ++ - common_params.causal_global ++ + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if cute.elem_less(key_pos, common_params.K) ++ and not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) + else: +- q_tok = common_params.blk_coord[1] +- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok +- tTR_rAcc[i] = ( +- tTR_rAcc[i] +- if cute.elem_less( +- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, +- k_bound, ++ mask_threshold = self.num_heads * ( ++ key_pos - common_params.K + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) +- else self.acc_dtype(-1.0e6) +- ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) + +@@ -2741,23 +2956,34 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if apply_mask: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): +- if cutlass.const_expr(self.fold_sq): +- q_tok = ( +- common_params.blk_coord[1] * self.fold_sq_ratio +- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) +- // self.num_heads ++ flat_q_row = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] ++ + common_params.blk_coord[0] * cta_m_rows ++ + tTR_tS[i][0] ++ ) ++ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index ++ if cutlass.const_expr(self.enable_dcp): ++ mask_threshold = self.num_heads * ( ++ self.cp_world * key_pos ++ + common_params.cp_rank ++ - common_params.causal_global ++ + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if cute.elem_less(key_pos, common_params.K) ++ and not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) + else: +- q_tok = common_params.blk_coord[1] +- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok +- tTR_rAcc[i] = ( +- tTR_rAcc[i] +- if cute.elem_less( +- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, +- k_bound, ++ mask_threshold = self.num_heads * ( ++ key_pos - common_params.K + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) +- else self.acc_dtype(-1.0e6) +- ) + # reduction for row_max after manual masking + row_max_new = tTR_rAcc.load().reduce( + cute.ReductionOp.MAX, row_max_new, 0 +@@ -3173,15 +3399,45 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + # load o + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + ++ row_has_key = True ++ if cutlass.const_expr(self.enable_dcp): ++ # The coordinate tensor returned by local_tile already carries ++ # the CTA's offset within the full M128 query tile. ++ flat_q_row = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] ++ + tTR_cO[0][0] ++ ) ++ first_key_threshold = self.num_heads * ( ++ self.cp_world * common_params.split_start_key ++ + common_params.cp_rank ++ - common_params.causal_global ++ + self.seq_len_q ++ ) ++ row_has_key = cute.elem_less( ++ common_params.split_start_key, common_params.K ++ ) and not cute.elem_less(flat_q_row, first_key_threshold) ++ + # apply output scale and normalize by row_sum +- for i in cutlass.range( +- cute.size(tTR_rAcc), vectorize=True, unroll_full=True +- ): +- tTR_rAcc[i] = ( +- tTR_rAcc[i] +- * epilogue_params.output_scale +- * cute.arch.rcp_approx(row_sum) ++ if cutlass.const_expr(self.enable_dcp): ++ normalization_scale = ( ++ epilogue_params.output_scale * cute.arch.rcp_approx(row_sum) ++ ) ++ normalization_scale = ( ++ normalization_scale if row_has_key else self.acc_dtype(0.0) + ) ++ for i in cutlass.range( ++ cute.size(tTR_rAcc), vectorize=True, unroll_full=True ++ ): ++ tTR_rAcc[i] = tTR_rAcc[i] * normalization_scale ++ else: ++ for i in cutlass.range( ++ cute.size(tTR_rAcc), vectorize=True, unroll_full=True ++ ): ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ * epilogue_params.output_scale ++ * cute.arch.rcp_approx(row_sum) ++ ) + + # store o to global memory + tR2G_rO_src = None +@@ -3262,6 +3518,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max + ) ++ if cutlass.const_expr(self.enable_dcp): ++ lse = lse if row_has_key else -self.lse_dtype.inf + # When writing directly to the user-facing mLSE (single-tile, + # no split-KV merge), convert from log2 base to natural log. + # When writing the per-split intermediate (mAccLSE branch), keep +@@ -3504,25 +3762,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + + return tile_sched_params, grid + +- @staticmethod +- def compute_fold_sq_ratio(num_heads: int, seq_len_q: int, m_tile: int) -> int: +- """Derive the seq_len_q-into-heads fold factor F. +- +- Returns the largest integer F such that: +- - F divides seq_len_q evenly +- - num_heads * F ≤ m_tile +- - 1 ≤ F ≤ seq_len_q +- +- F=1 means no folding (i.e. ``fold_sq`` should be False at the caller). +- """ +- if num_heads >= m_tile: +- return 1 +- max_fold = min(seq_len_q, m_tile // num_heads) +- for f in range(max_fold, 0, -1): +- if seq_len_q % f == 0: +- return f +- return 1 +- + @staticmethod + def get_workspace_size( + H: int, +@@ -3552,12 +3791,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + """ + if split_kv == 1: + return 0 +- # Decode packs heads into a physical 128-wide MMA-M tile. For H < 128, +- # split-KV partials can still touch the padded head lanes before +- # reduction, so size the workspace for max(H, 128). Mirrors the same +- # padding applied in initialize_workspace(). See #3235. +- workspace_heads = max(H, 128) +- return B * workspace_heads * S * split_kv * (D + 1) * acc_dtype.width // 8 ++ # The first workspace mode is the physical MMA-M row extent. Direct ++ # callers may still pass a logical H < 128, so preserve the M128 floor ++ # used by initialize_workspace(). The flat internal path passes 128. ++ workspace_rows = max(H, 128) ++ return B * workspace_rows * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( +@@ -3593,29 +3831,29 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): +- # Pad head dim to the physical 128-wide MMA-M tile. Without this, +- # H<128 split-KV partials write past the workspace. See #3235. +- workspace_H = cutlass.max(H, cutlass.Int32(128)) ++ # The first workspace mode is the physical MMA-M row extent. ++ # Preserve the M128 floor for direct callers that pass H < 128. ++ workspace_rows = cutlass.max(H, cutlass.Int32(128)) + align = 256 // self.q_dtype.width + acc_o_layout = cute.make_layout( +- (workspace_H, split_kv, D, S, B), ++ (workspace_rows, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, +- cute.assume(split_kv * workspace_H * D, align), +- cute.assume(workspace_H * split_kv * S * D, align), ++ cute.assume(split_kv * workspace_rows * D, align), ++ cute.assume(workspace_rows * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( +- (workspace_H, split_kv, S, B), ++ (workspace_rows, split_kv, S, B), + stride=( + split_kv, + 1, +- workspace_H * split_kv, +- workspace_H * split_kv * S, ++ workspace_rows * split_kv, ++ workspace_rows * split_kv * S, + ), + ) + acc_lse_iter = cute.recast_ptr( +@@ -3697,11 +3935,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP16: + return False + if is_var_split_kv and not is_var_seq: + return False +- if mma_qk_tiler_mn[0] < H: ++ if H <= 0 or mma_qk_tiler_mn[0] < H: + return False +- # When H < M tile, fold up to F tokens of S into H (M_eff = H*F ≤ M_tile). +- # F is auto-picked as the largest divisor of S with H*F ≤ M_tile. +- # F=1 always works, so any (H ≤ M_tile, S ≥ 1) is implementable. ++ # Query (token, head) rows are packed continuously into 128-row M ++ # tiles. A partial final tile is supported, so any H <= M_tile and ++ # S >= 1 is valid. + if S < 1: + return False + if K <= 0: +diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py +index f79683bc..f0aa6e72 100644 +--- a/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py ++++ b/flashinfer/cute_dsl/attention/monolithic/mla_decode_fp8.py +@@ -75,6 +75,7 @@ from cutlass.cutlass_dsl import BaseDSL + + from .mla_helpers import ( + ceil_div, ++ compute_q_tile_layout, + MAX_SPLITS, + LOG2_E, + MLAStaticTileScheduler, +@@ -106,15 +107,17 @@ launcher and ``flashinfer/cute_dsl/attention/mla_dispatch.py`` for impl selectio + + Constraints: + * Data type requirements: +- - Input/output: Float8E4M3FN ++ - Input: Float8E4M3FN ++ - Output: BFloat16 or Float8E4M3FN + - Accumulation and LSE: Float32 + * Fixed architecture parameters: +- - Number of attention heads: 128 ++ - Number of attention heads: 1-128 + - Latent dimension: 512 + - RoPE dimension: 64 + * Input query modes should be (NumHeads, LatentDim/RopeDim, SeqLenQ, BatchSize) + * Input kv latent/rope modes should be (SeqLenK, LatentDim/RopeDim, BatchSize) +-* Query sequence length must be 1-4 ++* Query sequence length must be positive; token/head rows are flattened across ++ 128-row M tiles with a safely padded final tile + * Only supports 2-CTA instructions + * Variable sequence length requires page table storage enabled + """ +@@ -136,7 +139,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + enable_pdl: bool, + num_heads: int = 128, + seq_len_q: int = 1, +- fold_sq: bool = False, ++ reducer_d_tiles: int = 1, ++ reducer_max_splits: int = MAX_SPLITS, ++ enable_dcp: bool = False, ++ cp_world: int = 1, + ): + """Initializes the configuration for a Blackwell Multi-Head Latent Attention (MLA) kernel. + +@@ -162,17 +168,24 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + :type is_var_split_kv: bool + :param enable_pdl: Whether to use PDL + :type enable_pdl: bool +- :param num_heads: Number of attention heads (pre-fold). Used for the +- per-row spec-decoding (MTP) causal mask q_token_index computation. ++ :param num_heads: Number of attention heads. Defines the flattened ++ ``(query_token, head)`` row geometry and division-free causal mask. + :type num_heads: int +- :param seq_len_q: Query sequence length (pre-fold). Combined with +- ``num_heads`` to derive the per-row q_token used by the causal mask. ++ :param seq_len_q: Query sequence length. Combined with ``num_heads`` ++ to define the flattened query-row extent and causal boundary. + :type seq_len_q: int +- :param fold_sq: Whether to fold tokens of ``seq_len_q`` into the head +- dimension so the M tile becomes [F sub_q_tok][num_heads heads]. +- Required when ``num_heads < mma_qk_tiler_mn[0]`` and ``seq_len_q > 1`` +- so the M tile is fully populated. +- :type fold_sq: bool ++ :param reducer_d_tiles: Number of independent D bands reduced per row. ++ :type reducer_d_tiles: int ++ :param reducer_max_splits: Compile-time reducer capacity. Direct users ++ retain the generic 256-split capacity by default; callers choosing ++ a smaller specialization must cap runtime split-KV accordingly. ++ :type reducer_max_splits: int ++ :param enable_dcp: Statically enable decode context-parallel causal ++ masking over a cyclic rank-local KV shard. ++ :type enable_dcp: bool ++ :param cp_world: Compile-time DCP world size. Rank-local key ``k`` maps ++ to global key ``k * cp_world + cp_rank``. ++ :type cp_world: int + """ + + self.latent_dim = 512 +@@ -187,25 +200,36 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + self.page_size = page_size + self.is_var_seq = is_var_seq + self.is_var_split_kv = is_var_split_kv ++ if not 1 <= reducer_max_splits <= MAX_SPLITS: ++ raise ValueError( ++ f"reducer_max_splits must be in [1, {MAX_SPLITS}], " ++ f"got {reducer_max_splits}" ++ ) ++ if is_var_split_kv and reducer_max_splits != MAX_SPLITS: ++ raise ValueError( ++ "variable split-KV requires the generic reducer capacity " ++ f"of {MAX_SPLITS}, got {reducer_max_splits}" ++ ) ++ self.reducer_max_splits = reducer_max_splits + self.enable_pdl = enable_pdl +- # Original (pre-fold) num_heads and seq_len_q used for per-row +- # spec-decoding (MTP) causal q_token_index computation. When fold_sq is +- # True the M tile is laid out as [F sub_q_tok][num_heads heads]; the +- # full q_tok for row r is blk_coord[1] * F + (r // num_heads). ++ if cp_world < 1: ++ raise ValueError(f"cp_world must be positive, got {cp_world}") ++ if not enable_dcp and cp_world != 1: ++ raise ValueError( ++ "cp_world must be 1 when decode context parallelism is disabled" ++ ) ++ self.enable_dcp = enable_dcp ++ self.cp_world = cp_world ++ # Flatten query-token and head modes into one affine row space. An MMA ++ # M tile may cross token boundaries; only the final tile can be ++ # partially populated and is padded by the Q TMA OOB fill. + self.num_heads = num_heads + self.seq_len_q = seq_len_q +- # fold_sq (caller-controlled): whether the folding code path is enabled. +- # fold_sq_ratio (derived): fold factor F ≥ 1; the largest divisor of +- # seq_len_q with num_heads * F ≤ M_tile and F ≤ seq_len_q. When the +- # caller passes fold_sq=False, the kernel ignores the ratio. +- # When fold_sq=True but the derived ratio is 1, the folding branch +- # is taken with F=1 (a no-op transform). +- self.fold_sq = fold_sq +- self.fold_sq_ratio = ( +- BlackwellMultiHeadLatentAttentionForwardFP8.compute_fold_sq_ratio( +- num_heads, seq_len_q, mma_qk_tiler_mn[0] +- ) +- ) ++ ( ++ self.total_q_rows, ++ self.num_q_tiles, ++ self.tail_q_rows, ++ ) = compute_q_tile_layout(num_heads, seq_len_q, mma_qk_tiler_mn[0]) + self.cluster_shape_mnk = (2, 1, 1) + self.use_2cta_instrs = True + # When using 2 CTAs with m=128: warps 0-1 handle accumulation for first half [0, n/2), +@@ -213,6 +237,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + self.warps_in_n = 2 + self.num_compute_warps = 4 + self.threads_per_warp = 32 ++ if reducer_d_tiles not in (1, 2, 4): ++ raise ValueError(f"unsupported reducer_d_tiles={reducer_d_tiles}") ++ self.reducer_d_tiles = reducer_d_tiles ++ self.reducer_d_tile = self.latent_dim // reducer_d_tiles + mma_qk_tiler_k = self.rope_dim * 2 + self.mma_qk_tiler = ( + self.mma_qk_tiler_mn[0], +@@ -294,13 +322,13 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + num_threads=(self.threads_per_warp * self.num_compute_warps), + ) + # Pingpong order barriers (OrderedSequenceBarrier<1,2> pattern). Each +- # group waits on its own bar via arrive_and_wait (= bar.sync); the OTHER +- # group signals via split-phase .arrive() (non-blocking). num_threads ++ # group waits on its own bar via arrive_and_wait (= bar.sync), while the ++ # other group normally signals via split-phase .arrive(). num_threads + # MUST cover both groups (256 = wait-side 128 + signal-side 128) so the + # bar releases only after BOTH have arrived — that's the cross-group +- # serialization that gives the TMEM peer-read its happens-before. +- # Init-phase trick: g1 pre-arrives bar_0 once at warp setup so g0's +- # first arrive_and_wait completes without waiting for g1's loop arrive. ++ # serialization that gives the TMEM peer-read its happens-before. The ++ # final owner also waits, closing the named-barrier generation before a ++ # persistent group can start its next logical work item. + self.softmax_order_bar_0 = pipeline.NamedBarrier( + barrier_id=5, + num_threads=(self.threads_per_warp * self.num_total_compute_warps), +@@ -353,6 +381,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + workspace: cute.Tensor, + split_kv: cutlass.Int32, + cache_seqs: Optional[cute.Tensor], ++ causal_seqlens_kv_global: Optional[cute.Tensor], ++ cp_rank: cutlass.Int32, + block_split_kvs: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, +@@ -388,6 +418,12 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + :type split_kv: cutlass.Int32 + :param cache_seqs: The cache sequences tensor with shape [batch_size] + :type cache_seqs: cute.Tensor ++ :param causal_seqlens_kv_global: Global exclusive causal bounds for the ++ newest query, with shape [batch_size]. Used only when DCP is ++ statically enabled. ++ :type causal_seqlens_kv_global: cute.Tensor ++ :param cp_rank: Runtime DCP rank in ``[0, cp_world)``. ++ :type cp_rank: cutlass.Int32 + :param block_split_kvs: The block split KV tensor with shape [batch_size] + :type block_split_kvs: cute.Tensor + :param softmax_scale: The scale factor for softmax +@@ -428,6 +464,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + q_latent = _reinterpret_4d(q_latent) + q_rope = _reinterpret_4d(q_rope) + o = _reinterpret_4d(o) ++ # Keep unpacked output aliases for the split-KV reducer, which ++ # launches only real (head, q_token) rows. ++ o_unpacked = o + + # Reinterpret contiguous [num_pages, page_size, D] as [page_size, D, num_pages] + # Input stride: (PS*D, D, 1) → Target: (D, 1, PS*D) +@@ -461,54 +500,53 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + stride=(lse.stride[2], lse.stride[1], lse.stride[0]), + ), + ) ++ lse_unpacked = lse + +- # When num_heads < M tile, fold up to F = fold_sq_ratio tokens of +- # seq_len_q into the head dimension so M_eff = num_heads * F (≤ M_tile). +- # E.g., H=32, S_q=4 → F=4, M_eff=128, S_q_eff=1 +- # E.g., H=32, S_q=8 → F=4, M_eff=128, S_q_eff=2 +- # This works because MLA shares KV across all heads/queries independently. +- # Tensor layout: [H, D, S_q, B] → [H*F, D, S_q/F, B]; relies on +- # stride_S == stride_H * H (always true for contiguous [B, S_q, H, D] +- # tensors after _reinterpret_4d). +- if cutlass.const_expr(self.fold_sq): +- F = self.fold_sq_ratio +- +- def _fold_sq_4d(t): +- return cute.make_tensor( +- t.iterator, +- cute.make_layout( +- ( +- t.shape[0] * F, +- t.shape[1], +- t.shape[2] // F, +- t.shape[3], +- ), +- stride=( +- t.stride[0], +- t.stride[1], +- t.stride[2] * F, +- t.stride[3], +- ), +- ), +- ) +- +- q_latent = _fold_sq_4d(q_latent) +- q_rope = _fold_sq_4d(q_rope) +- o = _fold_sq_4d(o) +- # LSE: [H, S_q, B] → [H*F, S_q/F, B] +- lse = cute.make_tensor( +- lse.iterator, ++ # Flatten [H, D, S_q, B] into the affine row order ++ # flat_row = q_token * H + q_head. A single bounded TMA descriptor can ++ # then cross token boundaries and OOB-zero-fill only the final M tile. ++ def _flatten_q_rows(t): ++ return cute.make_tensor( ++ t.iterator, + cute.make_layout( +- (lse.shape[0] * F, lse.shape[1] // F, lse.shape[2]), +- stride=(lse.stride[0], lse.stride[1] * F, lse.stride[2]), ++ (self.total_q_rows, t.shape[1], t.shape[3]), ++ stride=(t.stride[0], t.stride[1], t.stride[3]), + ), + ) + ++ q_latent = _flatten_q_rows(q_latent) ++ q_rope = _flatten_q_rows(q_rope) ++ ++ # O/LSE expose one physical M tile per query tile. Their final tile's ++ # virtual padding is never dereferenced because stores are ++ # predicated by get_valid_q_rows(). ++ m_tile = self.mma_qk_tiler_mn[0] ++ runtime_num_q_tiles = cute.ceil_div(o.shape[0] * o.shape[2], m_tile) ++ o = cute.make_tensor( ++ o.iterator, ++ cute.make_layout( ++ (m_tile, o.shape[1], runtime_num_q_tiles, o.shape[3]), ++ stride=( ++ o.stride[0], ++ o.stride[1], ++ o.stride[0] * m_tile, ++ o.stride[3], ++ ), ++ ), ++ ) ++ lse = cute.make_tensor( ++ lse.iterator, ++ cute.make_layout( ++ (m_tile, runtime_num_q_tiles, lse.shape[2]), ++ stride=(lse.stride[0], lse.stride[0] * m_tile, lse.stride[2]), ++ ), ++ ) ++ + acc_o, acc_lse = self.initialize_workspace( +- q_latent.shape[0], +- q_latent.shape[1], +- q_latent.shape[2], +- q_latent.shape[3], ++ o.shape[0], ++ o.shape[1], ++ o.shape[2], ++ o.shape[3], + split_kv, + self.acc_dtype, + workspace, +@@ -841,6 +879,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + acc_lse, + split_kv, + cache_seqs, ++ causal_seqlens_kv_global, ++ cp_rank, + block_split_kvs, + softmax_scale_log2, + output_scale, +@@ -867,17 +907,21 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + ) + if cutlass.const_expr(acc_o is not None): + self.reduction_kernel( +- o, +- lse, ++ o_unpacked, ++ lse_unpacked, + acc_o, + acc_lse, + split_kv, + cache_seqs, + block_split_kvs, + ).launch( +- grid=(q_latent.shape[0], q_latent.shape[2], q_latent.shape[3]), ++ grid=( ++ o_unpacked.shape[0] * self.reducer_d_tiles, ++ o_unpacked.shape[2], ++ o_unpacked.shape[3], ++ ), + block=[self.threads_per_warp * self.num_compute_warps, 1, 1], +- smem=MAX_SPLITS * self.acc_dtype.width // 8, ++ smem=self.reducer_max_splits * self.acc_dtype.width // 8, + stream=stream, + min_blocks_per_mp=1, + use_pdl=self.enable_pdl, +@@ -943,6 +987,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + mAccLSE: Optional[cute.Tensor], + split_kv: cutlass.Int32, + cache_seqs: cute.Tensor, ++ causal_seqlens_kv_global: Optional[cute.Tensor], ++ cp_rank: cutlass.Int32, + block_split_kvs: cute.Tensor, + softmax_scale_log2: cutlass.Float32, + output_scale: cutlass.Float32, +@@ -985,8 +1031,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + :type mQL: cute.Tensor + :param tma_atom_q_rope: TMA copy atom for query rope tensor + :type tma_atom_q_rope: cute.CopyAtom +- :param mKR: Compressed rope tensor +- :type mKR: cute.Tensor ++ :param mQR: query rope tensor ++ :type mQR: cute.Tensor + :param tma_atom_c_latent: TMA copy atom for c latent tensor + :type tma_atom_c_latent: cute.CopyAtom + :param mCL: Compressed latent tensor +@@ -1009,6 +1055,12 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + :type split_kv: cutlass.Int32 + :param cache_seqs: The variable sequence length tensor + :type cache_seqs: cute.Tensor ++ :param causal_seqlens_kv_global: Global exclusive causal bounds for the ++ newest query. Present only in the statically enabled DCP ++ specialization. ++ :type causal_seqlens_kv_global: cute.Tensor ++ :param cp_rank: Runtime rank of this cyclic KV shard ++ :type cp_rank: cutlass.Int32 + :param block_split_kvs: The per-block split_kv values tensor + :type block_split_kvs: cute.Tensor + :param softmax_scale_log2: The log2 scale factor for softmax +@@ -1425,6 +1477,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: ++ causal_seq_len = cutlass.Int32(0) ++ if cutlass.const_expr(self.enable_dcp): ++ causal_seq_len = causal_seqlens_kv_global[blk_coord[2]] + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, +@@ -1433,6 +1488,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], ++ causal_seq_len=causal_seq_len, ++ cp_rank=cp_rank, + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, +@@ -1492,17 +1549,18 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + mma_s_consumer_state.advance() + p_mma_producer_state.advance() + p_cor_producer_state.advance() +- # Pingpong init-phase trick: g1 pre-arrives softmax_order_bar_0 once +- # so g0's FIRST softmax_order_bar_0.arrive_and_wait() completes +- # without waiting for g1's first loop arrive. Replaces a per-iter +- # is_first_tile guard. Bar is num_total_compute_warps (256 threads) +- # so wait-side 128 + signal-side 128 = 256 → release. ++ # Each nonempty compute() call initializes g1's peer metadata and ++ # pre-arrives softmax_order_bar_0 for that logical work item. This ++ # lets g0's first wait complete without a first-tile special case. + while work_tile.is_valid_tile: + blk_coord = work_tile.tile_idx + k_index, k_tile_count, local_split_kv = self.get_k_tile_count( + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: ++ causal_seq_len = cutlass.Int32(0) ++ if cutlass.const_expr(self.enable_dcp): ++ causal_seq_len = causal_seqlens_kv_global[blk_coord[2]] + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, +@@ -1511,6 +1569,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], ++ causal_seq_len=causal_seq_len, ++ cp_rank=cp_rank, + L=mCL.shape[1], + tmem_ptr=tmem_ptr, + tidx=tidx, +@@ -1568,16 +1628,22 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + split_kv, cache_seqs, block_split_kvs, blk_coord + ) + if k_tile_count > 0: ++ causal_seq_len = cutlass.Int32(0) ++ if cutlass.const_expr(self.enable_dcp): ++ causal_seq_len = causal_seqlens_kv_global[blk_coord[2]] + compute_common_params = SimpleNamespace( + blk_coord=blk_coord, + split_kv=split_kv, + local_split_kv=local_split_kv, ++ k_index=k_index, + smem_exchange=epilogue_smem_exchange, + mAccO=mAccO, + mO=mO, + K=cache_seqs[blk_coord[2]], ++ causal_seq_len=causal_seq_len, ++ cp_rank=cp_rank, + L=mCL.shape[1], +- H=mQL.shape[0], ++ H=self.get_valid_q_rows(blk_coord[1]), + tmem_ptr=tmem_ptr, + tidx=tidx, + tiled_mma_pv=tiled_mma_pv, +@@ -1597,10 +1663,73 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + p_cor_consumer_state=p_cor_consumer_state, + mma_o_consumer_state=mma_o_consumer_state, + ) ++ elif cutlass.const_expr(self.enable_dcp): ++ self.store_empty_dcp_work( ++ mO, ++ mLSE, ++ mAccO, ++ mAccLSE, ++ blk_coord, ++ self.get_valid_q_rows(blk_coord[1]), ++ tidx, ++ ) + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + return + ++ @cute.jit ++ def store_empty_dcp_work( ++ self, ++ mO: Optional[cute.Tensor], ++ mLSE: Optional[cute.Tensor], ++ mAccO: Optional[cute.Tensor], ++ mAccLSE: Optional[cute.Tensor], ++ blk_coord: cute.Coord, ++ valid_q_rows: cutlass.Int32, ++ tidx: cutlass.Int32, ++ ): ++ """Materialize O=0 and LSE=-inf for a DCP split with no physical K tile.""" ++ cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] ++ compute_threads = self.num_compute_warps * self.threads_per_warp ++ local_tidx = tidx % compute_threads ++ cta_row_base = blk_coord[0] * cta_m_rows ++ ++ for linear_idx in cutlass.range( ++ local_tidx, cta_m_rows * self.latent_dim, compute_threads ++ ): ++ cta_row = linear_idx // self.latent_dim ++ d_idx = linear_idx % self.latent_dim ++ q_row = cta_row_base + cta_row ++ if cute.elem_less(q_row, valid_q_rows): ++ if cutlass.const_expr(mAccO is None): ++ mO[q_row, d_idx, blk_coord[1], blk_coord[2]] = self.o_dtype(0.0) ++ else: ++ mAccO[ ++ q_row, ++ blk_coord[3], ++ d_idx, ++ blk_coord[1], ++ blk_coord[2], ++ ] = self.acc_dtype(0.0) ++ ++ if cute.elem_less(local_tidx, cta_m_rows): ++ q_row = cta_row_base + local_tidx ++ if cute.elem_less(q_row, valid_q_rows): ++ if cutlass.const_expr(mAccLSE is None): ++ mLSE[q_row, blk_coord[1], blk_coord[2]] = -self.lse_dtype.inf ++ else: ++ mAccLSE[ ++ q_row, blk_coord[3], blk_coord[1], blk_coord[2] ++ ] = -self.lse_dtype.inf ++ ++ @cute.jit ++ def get_valid_q_rows(self, q_tile_idx: cutlass.Int32) -> cutlass.Int32: ++ """Number of valid flattened query rows in one physical M tile.""" ++ valid_rows = self.mma_qk_tiler_mn[0] ++ if q_tile_idx == self.num_q_tiles - 1: ++ valid_rows = self.tail_q_rows ++ return valid_rows ++ + @cute.kernel + def reduction_kernel( + self, +@@ -1632,27 +1761,47 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + """ + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() +- blk_coord = (bidx, bidy, bidz) ++ # Reducer blocks cover one D band of one real output row. Splitting ++ # D512 into independent D128 bands exposes enough CTAs for small-batch ++ # decode without any cross-CTA reduction or synchronization. ++ d_tile_idx = bidx % self.reducer_d_tiles ++ blk_coord = (bidx // self.reducer_d_tiles, bidy, bidz) ++ flat_q_row = blk_coord[1] * self.num_heads + blk_coord[0] ++ # The physical M tile is fixed at 128 rows, so map with shift/mask. ++ q_tile = flat_q_row >> 7 ++ q_tile_row = flat_q_row & 127 + local_split_kv = ( + block_split_kvs[blk_coord[2]] if self.is_var_split_kv else split_kv + ) + k_tile_total = cute.ceil_div(cache_seqs[blk_coord[2]], self.mma_qk_tiler[1]) +- k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) +- local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) ++ if cutlass.const_expr(self.enable_dcp): ++ k_tile_per_cta = cutlass.max( ++ cute.ceil_div(k_tile_total, local_split_kv), cutlass.Int32(1) ++ ) ++ local_split_kv = cutlass.max( ++ cute.ceil_div(k_tile_total, k_tile_per_cta), cutlass.Int32(1) ++ ) ++ else: ++ k_tile_per_cta = cute.ceil_div(k_tile_total, local_split_kv) ++ local_split_kv = cute.ceil_div(k_tile_total, k_tile_per_cta) + + # Alloc shared memory + smem = utils.SmemAllocator() +- storage = smem.allocate(MAX_SPLITS * self.acc_dtype.width // 8, 16) ++ storage = smem.allocate(self.reducer_max_splits * self.acc_dtype.width // 8, 16) + lse_scale_ptr = cute.recast_ptr(storage, dtype=self.acc_dtype) +- smem_lse_scale = cute.make_tensor(lse_scale_ptr, cute.make_layout(MAX_SPLITS)) ++ smem_lse_scale = cute.make_tensor( ++ lse_scale_ptr, cute.make_layout(self.reducer_max_splits) ++ ) + + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_wait() +- gLSE = mAccLSE[blk_coord[0], None, blk_coord[1], blk_coord[2]] ++ gLSE = mAccLSE[q_tile_row, None, q_tile, blk_coord[2]] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + # calculate the global lse and exp ^ (local_lse - global_lse) +- lse_per_thread = cute.ceil_div(MAX_SPLITS, self.threads_per_warp) ++ lse_per_thread = cute.ceil_div( ++ self.reducer_max_splits, self.threads_per_warp ++ ) + + local_lse = cute.make_rmem_tensor( + cute.make_layout(lse_per_thread), self.lse_dtype +@@ -1669,39 +1818,59 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + # reduce the local lse + lse_max = cute.arch.fmax(lse_max, local_lse[i]) + lse_max = cute.arch.warp_reduction_max(lse_max) +- lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 ++ if cutlass.const_expr(self.enable_dcp): ++ has_valid_lse = lse_max != -self.lse_dtype.inf ++ lse_max = lse_max if has_valid_lse else 0.0 ++ else: ++ lse_max = lse_max if lse_max != -self.lse_dtype.inf else 0.0 + # calculate sum_lse + sum_lse = 0.0 + for i in cutlass.range_constexpr(lse_per_thread): + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + sum_lse = cute.arch.warp_reduction_sum(sum_lse) + # calculate the global_lse +- global_lse = ( +- lse_max + cute.math.log2(sum_lse, fastmath=True) +- if not sum_lse == self.lse_dtype(0.0) or sum_lse != sum_lse # noqa: SIM201 +- else self.lse_dtype.inf +- ) +- if tidx == 0: +- # Convert from kernel-internal log2 base to the natural-log +- # convention exposed to callers (matches trtllm-gen / flash-attn). +- # `1.0 / LOG2_E == ln(2)`. +- mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( +- 1.0 / LOG2_E ++ if cutlass.const_expr(self.enable_dcp): ++ global_lse = ( ++ lse_max + cute.math.log2(sum_lse, fastmath=True) ++ if has_valid_lse ++ else -self.lse_dtype.inf ++ ) ++ else: ++ global_lse = ( ++ lse_max + cute.math.log2(sum_lse, fastmath=True) ++ if not sum_lse == self.lse_dtype(0.0) # noqa: SIM201 ++ or sum_lse != sum_lse ++ else self.lse_dtype.inf + ) ++ if d_tile_idx == 0: ++ if tidx == 0: ++ # Convert from kernel-internal log2 base to the natural-log ++ # convention exposed to callers (matches trtllm-gen / flash-attn). ++ # `1.0 / LOG2_E == ln(2)`. ++ mLSE[blk_coord[0], blk_coord[1], blk_coord[2]] = global_lse * ( ++ 1.0 / LOG2_E ++ ) + # store the scale to shared memory + for i in cutlass.range_constexpr(lse_per_thread): + split_kv_idx = tidx + i * self.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): +- smem_lse_scale[split_kv_idx] = cute.math.exp2( +- local_lse[i] - global_lse, fastmath=True +- ) ++ if cutlass.const_expr(self.enable_dcp): ++ smem_lse_scale[split_kv_idx] = ( ++ cute.math.exp2(local_lse[i] - global_lse, fastmath=True) ++ if has_valid_lse ++ else self.acc_dtype(0.0) ++ ) ++ else: ++ smem_lse_scale[split_kv_idx] = cute.math.exp2( ++ local_lse[i] - global_lse, fastmath=True ++ ) + + pipeline.sync(barrier_id=4) + + elements_per_thread = cute.ceil_div( +- self.latent_dim, self.threads_per_warp * self.num_compute_warps ++ self.reducer_d_tile, self.threads_per_warp * self.num_compute_warps + ) +- gAccO = mAccO[blk_coord[0], None, None, blk_coord[1], blk_coord[2]] ++ gAccO = mAccO[q_tile_row, None, None, q_tile, blk_coord[2]] + rAccO = cute.make_rmem_tensor( + cute.make_layout(elements_per_thread), self.acc_dtype + ) +@@ -1709,11 +1878,19 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + rAccO.fill(0.0) + for i in range(local_split_kv): + for j in cutlass.range_constexpr(elements_per_thread): +- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps ++ element_idx = ( ++ d_tile_idx * self.reducer_d_tile ++ + tidx ++ + j * self.threads_per_warp * self.num_compute_warps ++ ) + rAccO[j] += gAccO[i, element_idx] * smem_lse_scale[i] + rO.store(rAccO.load().to(self.o_dtype)) + for j in cutlass.range_constexpr(elements_per_thread): +- element_idx = tidx + j * self.threads_per_warp * self.num_compute_warps ++ element_idx = ( ++ d_tile_idx * self.reducer_d_tile ++ + tidx ++ + j * self.threads_per_warp * self.num_compute_warps ++ ) + mO[blk_coord[0], element_idx, blk_coord[1], blk_coord[2]] = rO[j] + if cutlass.const_expr(self.enable_pdl): + cute.arch.griddepcontrol_launch_dependents() +@@ -1820,7 +1997,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + # page table + mPT = common_params.mPT[None, common_params.blk_coord[2]] + +- # Flatten divide and partition global tensors for QK TMA load ++ # The descriptor has one bounded flat M mode, so TMA can cross query ++ # token boundaries and zero-fills only the final physical tile. + # (bM, bK, rM, rK, rL) + mma_qk_tiler_mk = cute.select(self.mma_qk_tiler, mode=[0, 2]) + gQL = cute.flat_divide(qk_params.mQL, mma_qk_tiler_mk) +@@ -1898,11 +2076,14 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + tSgKR, + ) + ++ # Q is flattened as (Sq * Hq, D, B). The logical query-tile ++ # coordinate therefore lives in the residual-M mode created by ++ # flat_divide, rather than in a separate grouped-query mode. + tQLgQL = tQLgQL_mkl[ +- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] ++ None, common_params.blk_coord[1], None, common_params.blk_coord[2] + ] + tQRgQR = tQRgQR_mkl[ +- None, None, None, common_params.blk_coord[1], common_params.blk_coord[2] ++ None, common_params.blk_coord[1], None, common_params.blk_coord[2] + ] + + # set extra params +@@ -2002,6 +2183,33 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + k_tile_count -= 1 + return load_v_producer_state + ++ @cute.jit ++ def issue_q_tma_load( ++ self, ++ tma_atom_q_latent: cute.CopyAtom, ++ tQLgQL: cute.Tensor, ++ tQsQ: cute.Tensor, ++ tma_atom_q_rope: cute.CopyAtom, ++ tQRgQR: cute.Tensor, ++ tQsQ_rope: cute.Tensor, ++ tma_bar_ptr, ++ ): ++ """Issue one bounded flattened query-tile load into shared Q stages.""" ++ for i in cutlass.range_constexpr(self.iterations_qk_latent): ++ cute.copy( ++ tma_atom_q_latent, ++ tQLgQL[None, i], ++ tQsQ[None, (i, 0)], ++ tma_bar_ptr=tma_bar_ptr, ++ ) ++ for i in cutlass.range_constexpr(self.iterations_qk_rope): ++ cute.copy( ++ tma_atom_q_rope, ++ tQRgQR[None, i], ++ tQsQ_rope[None, i], ++ tma_bar_ptr=tma_bar_ptr, ++ ) ++ + @cute.jit + def load_tma_qk_one_k_tile( + self, +@@ -2057,22 +2265,15 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + tma_bar_ptr = load_q_pipeline.producer_get_barrier(load_q_producer_state) + # expect the extra bytes for q. + load_q_pipeline.producer_acquire(load_q_producer_state) +- for i in cutlass.range_constexpr(self.iterations_qk_latent): +- # load q latent +- cute.copy( +- qk_params.tma_atom_q_latent, +- qk_params.tQLgQL[None, 0, i], +- qk_params.tQsQ[None, (i, 0)], +- tma_bar_ptr=tma_bar_ptr, +- ) +- for i in cutlass.range_constexpr(self.iterations_qk_rope): +- # load q rope +- cute.copy( +- qk_params.tma_atom_q_rope, +- qk_params.tQRgQR[None, 0, i], +- qk_params.tQsQ_rope[None, i], +- tma_bar_ptr=tma_bar_ptr, +- ) ++ self.issue_q_tma_load( ++ qk_params.tma_atom_q_latent, ++ qk_params.tQLgQL, ++ qk_params.tQsQ, ++ qk_params.tma_atom_q_rope, ++ qk_params.tQRgQR, ++ qk_params.tQsQ_rope, ++ tma_bar_ptr, ++ ) + load_q_producer_state.advance() + # get the mbar ptr from pipeline. + tma_bar_ptr = common_params.load_k_pipeline.producer_get_barrier( +@@ -2525,7 +2726,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + row_max = self.acc_dtype(self.init_row_max) + row_sum = self.acc_dtype(0) + correction_factor = self.acc_dtype(1) +- odd_k_tile = k_tile_count % 2 == 1 ++ k_tile_count_init = k_tile_count ++ odd_k_tile = k_tile_count_init % 2 == 1 + # 2softmax: g0 takes even k-tiles, g1 takes odd k-tiles. g1 advances all + # pipeline states once at entry to start on stage 1 (one stage per group). + if cutlass.const_expr(is_second_compute_warp): +@@ -2544,26 +2746,59 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + common_params, softmax_params, p_cor_producer_state + ) + self.softmax_order_bar_0.arrive() +- # Number of tiles from the global-K end that may contain causal-masked +- # positions. Min k_bound = K - (S_q-1), which can span up to +- # ceil((seq_len_q-2)/tile_N)+1 tiles (tile-boundary-crossing case). For +- # S_q=1 this reduces to 1 tile — identical to a plain K-bound check. ++ # The first tile that can contain a key at or beyond this query tile's ++ # earliest causal bound. A flat M tile starts at row q_tile * M, so ++ # floor(row / H) is its earliest query token. Later query tiles can ++ # therefore keep more fully dense K tiles on the compile-time unmasked ++ # path. The token division is outside the per-score loop (and is ++ # replicated by each participating CTA/compute group); tile_n is a ++ # static power of two, so the remaining division is a shift. + tile_n = self.mma_qk_tiler[1] +- mask_tile_count = (self.seq_len_q - 2 + tile_n - 1) // tile_n + 1 ++ first_q_token = cutlass.Int32(0) ++ if cutlass.const_expr(self.num_q_tiles > 1): ++ first_q_token = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] ++ ) // self.num_heads ++ if cutlass.const_expr(self.enable_dcp): ++ # The earliest chronological query represented by this physical ++ # M128 tile has the smallest visible rank-local K bound. Compute ++ # that bound once per query tile, then cap it by the physical ++ # rank-local K extent so a partial physical tail is still masked. ++ dcp_bound_numerator = cutlass.max( ++ common_params.causal_seq_len ++ - common_params.cp_rank ++ - (self.seq_len_q - 1) ++ + first_q_token, ++ cutlass.Int32(0), ++ ) ++ earliest_local_bound = ( ++ dcp_bound_numerator + self.cp_world - 1 ++ ) // self.cp_world ++ effective_local_bound = cutlass.min(common_params.K, earliest_local_bound) ++ first_mask_tile_idx = effective_local_bound // tile_n ++ else: ++ first_mask_tile_idx = cutlass.max( ++ (common_params.K - self.seq_len_q + 1 + first_q_token) // tile_n, ++ cutlass.Int32(0), ++ ) + +- # first_mask_tile_idx is the global index of the first tile that may +- # need masking. Runtime because it depends on K (per-batch in +- # var-seq / split-KV). +- first_mask_tile_idx = k_tile_total - mask_tile_count ++ # The non-split two-softmax path still needs its global final tile in ++ # phase 2 for final metadata exchange, even when that tile is dense. ++ # Masking itself remains controlled by first_mask_tile_idx below. ++ first_phase_2_tile_idx = cutlass.min(first_mask_tile_idx, k_tile_total - 1) + + # Phase 1: pure unmasked bulk tiles (all columns strictly < min k_bound). + # 2softmax: each group steps by 2 k-tiles; phase boundaries respect this. +- while k_tile_count > 0 and k_index < first_mask_tile_idx: ++ while k_tile_count > 0 and k_index < first_phase_2_tile_idx: + is_local_last_tile = ( + False + if cutlass.const_expr(common_params.mAccO is None) + else k_tile_count == 1 + ) ++ if cutlass.const_expr(is_second_compute_warp): ++ is_order_chain_last_tile = k_tile_count == 1 and not odd_k_tile ++ else: ++ is_order_chain_last_tile = k_tile_count == 1 and odd_k_tile + ( + mma_s_consumer_state, + p_mma_producer_state, +@@ -2584,6 +2819,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + is_second_compute_warp, + False, + is_local_last_tile, ++ is_order_chain_last_tile, + ) + k_index = k_index + 2 + k_tile_count = k_tile_count - 1 +@@ -2600,6 +2836,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + # Phase 2: remaining tiles that overlap the causal / K-bound region, + # including this work-split's final tile. + while k_tile_count > 0: ++ if cutlass.const_expr(is_second_compute_warp): ++ is_order_chain_last_tile = k_tile_count == 1 and not odd_k_tile ++ else: ++ is_order_chain_last_tile = k_tile_count == 1 and odd_k_tile + ( + mma_s_consumer_state, + p_mma_producer_state, +@@ -2618,8 +2858,9 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + row_sum, + correction_factor, + is_second_compute_warp, ++ k_index >= first_mask_tile_idx, + True, +- True, ++ is_order_chain_last_tile, + ) + k_index = k_index + 2 + k_tile_count = k_tile_count - 1 +@@ -2635,14 +2876,17 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + + if odd_k_tile: + if cutlass.const_expr(is_second_compute_warp): +- # next first compute warp in this wave +- p_mma_producer_state.advance() +- mma_s_consumer_state.advance() +- p_cor_producer_state.advance() +- # first compute warp in next wave ++ # Keep g1 exactly one pipeline entry ahead of g0 at the next ++ # logical work item. When this split has one tile, g1 owns no ++ # tile and needs only this single advance. For larger odd ++ # splits it also skips the final g0-owned stage below. + p_mma_producer_state.advance() + mma_s_consumer_state.advance() + p_cor_producer_state.advance() ++ if k_tile_count_init > 1: ++ p_mma_producer_state.advance() ++ mma_s_consumer_state.advance() ++ p_cor_producer_state.advance() + else: + p_mma_producer_state.advance() + mma_s_consumer_state.advance() +@@ -2949,6 +3193,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + is_second_compute_warp: bool, + apply_mask: bool, + is_local_last_tile: cutlass.Boolean, ++ is_order_chain_last_tile: cutlass.Boolean, + ) -> tuple[ + pipeline.PipelineState, + pipeline.PipelineState, +@@ -2978,11 +3223,14 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + :param correction_factor: The correction factor + :type correction_factor: cutlass.Float32 + :param apply_mask: Whether the tile needs K-bound / causal masking (Python bool +- for the unmasked/masked bulk loops; runtime cutlass.Boolean for the +- split-KV final iter where mask only applies on the global last tile). ++ for the unmasked/masked bulk loops; runtime cutlass.Boolean for a ++ work-split final tile at the exact per-query-tile mask boundary). + :type apply_mask: bool | cutlass.Boolean + :param is_local_last_tile: Whether the last tile is local + :type is_local_last_tile: cutlass.Boolean ++ :param is_order_chain_last_tile: Whether this group owns the final tile ++ in the logical work item's cross-group order chain ++ :type is_order_chain_last_tile: cutlass.Boolean + + :return: The MMA s consumer state, the P MMA producer state, the P correction producer state, the row max, the row sum, and the correction factor + :rtype: tuple[pipeline.PipelineState, pipeline.PipelineState, pipeline.PipelineState, cutlass.Float32, cutlass.Float32, cutlass.Float32] +@@ -3028,40 +3276,52 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + tTR_rAcc = cute.make_fragment_like(tTR_tS, self.acc_dtype) + + row_max_new = row_max +- # Spec-decoding (MTP) causal mask: each row represents one (q_token, head) +- # pair; row r's effective K bound is K - (S_q - 1 - q_tok(r)). +- # With fold factor F = self.fold_sq_ratio (fold_sq=True), the M tile is +- # laid out as [F sub_q_tok][num_heads heads] and there are S_q/F outer +- # chunks indexed by blk_coord[1]: +- # q_tok(r) = blk_coord[1] * F + (r_global // num_heads) +- # r_global = row_in_cta + cluster_idx * (M_tile / cluster_m) +- # When fold_sq=False this reduces to q_tok = blk_coord[1]. For S_q=1 +- # this further reduces to k_bound = K (plain K-bound check). ++ # Spec-decoding (MTP) causal mask. For flattened row ++ # r = q_token * H + q_head, ++ # the usual key < K - S_q + 1 + q_token predicate is equivalent to ++ # H * (key - K + S_q) <= r. ++ # With cyclic DCP, local key k maps to W*k+rank, producing ++ # H * (W*k + rank - G + S_q) <= r. ++ # Both forms avoid integer division/modulo. The enclosing apply_mask ++ # branch remains compile-time false for dense bulk K tiles, so they do ++ # not execute any of this row-dependent arithmetic. + # Masked positions are filled with a large negative sentinel (not -inf) +- # to avoid NaN propagation when an entire row becomes masked. ++ # to avoid NaN propagation when an entire row becomes masked. The DCP ++ # epilogue turns such an empty row into the neutral O=0, LSE=-inf pair. + cta_m_rows = self.mma_qk_tiler[0] // self.cluster_shape_mnk[0] + arch = BaseDSL._get_dsl().get_arch_enum() + if cutlass.const_expr(arch >= Arch.sm_100 and arch <= Arch.sm_100f): + cute.copy(tmem_tiled_copy, tTR_tAcc, tTR_rAcc) + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): + if apply_mask: +- if cutlass.const_expr(self.fold_sq): +- q_tok = ( +- common_params.blk_coord[1] * self.fold_sq_ratio +- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) +- // self.num_heads ++ flat_q_row = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] ++ + common_params.blk_coord[0] * cta_m_rows ++ + tTR_tS[i][0] ++ ) ++ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index ++ if cutlass.const_expr(self.enable_dcp): ++ mask_threshold = self.num_heads * ( ++ self.cp_world * key_pos ++ + common_params.cp_rank ++ - common_params.causal_seq_len ++ + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if cute.elem_less(key_pos, common_params.K) ++ and not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) + else: +- q_tok = common_params.blk_coord[1] +- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok +- tTR_rAcc[i] = ( +- tTR_rAcc[i] +- if cute.elem_less( +- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, +- k_bound, ++ mask_threshold = self.num_heads * ( ++ key_pos - common_params.K + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) +- else self.acc_dtype(-1.0e6) +- ) + # reduction for row_max + row_max_new = tTR_rAcc.load().reduce(cute.ReductionOp.MAX, row_max_new, 0) + elif cutlass.const_expr( +@@ -3092,23 +3352,34 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + tTR_rAcc = cute.make_tensor(tTR_rAcc_red.iterator, tTR_rAcc.layout) + if apply_mask: + for i in cutlass.range_constexpr(cute.size(tTR_rAcc)): +- if cutlass.const_expr(self.fold_sq): +- q_tok = ( +- common_params.blk_coord[1] * self.fold_sq_ratio +- + (tTR_tS[i][0] + common_params.blk_coord[0] * cta_m_rows) +- // self.num_heads ++ flat_q_row = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] ++ + common_params.blk_coord[0] * cta_m_rows ++ + tTR_tS[i][0] ++ ) ++ key_pos = tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index ++ if cutlass.const_expr(self.enable_dcp): ++ mask_threshold = self.num_heads * ( ++ self.cp_world * key_pos ++ + common_params.cp_rank ++ - common_params.causal_seq_len ++ + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if cute.elem_less(key_pos, common_params.K) ++ and not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) + else: +- q_tok = common_params.blk_coord[1] +- k_bound = common_params.K - (self.seq_len_q - 1) + q_tok +- tTR_rAcc[i] = ( +- tTR_rAcc[i] +- if cute.elem_less( +- tTR_tS[i][1] + self.mma_qk_tiler[1] * k_index, +- k_bound, ++ mask_threshold = self.num_heads * ( ++ key_pos - common_params.K + self.seq_len_q ++ ) ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ if not cute.elem_less(flat_q_row, mask_threshold) ++ else self.acc_dtype(-1.0e6) + ) +- else self.acc_dtype(-1.0e6) +- ) + # reduction for row_max after manual masking + row_max_new = tTR_rAcc.load().reduce( + cute.ReductionOp.MAX, row_max_new, 0 +@@ -3152,7 +3423,7 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + # the .arrive() it issues at the bottom of its prev iter, so this + # arrive_and_wait gives us the acquire memory ordering needed for the + # load_other_group_metadata read below. g0's first wait is satisfied +- # by g1's pre-arrive of bar_0 at warp setup (init-phase trick). ++ # by g1's per-work pre-arrive of bar_0 (init-phase trick). + if cutlass.const_expr(is_second_compute_warp): + self.softmax_order_bar_1.arrive_and_wait() + else: +@@ -3285,9 +3556,10 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + ) + + # Pingpong B — signal the OTHER group to start its critical section. +- # Split-phase .arrive() contributes this group's threads to the other +- # group's bar and falls through immediately. cfence: control-flow fence +- # to prevent ptxas from hoisting subsequent work above bar.arrive. ++ # Intermediate tiles use split-phase arrive so the groups overlap. The ++ # final owner waits for the peer's parity-tail rendezvous, preventing a ++ # persistent group from arriving on the same named-barrier generation ++ # again in the next logical work item before this one has reset. + + # split kv case + if is_local_last_tile: +@@ -3304,9 +3576,15 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + ) + # cute.nvgpu.cfence() + if cutlass.const_expr(is_second_compute_warp): +- self.softmax_order_bar_0.arrive() # g1 → g0 ++ if is_order_chain_last_tile: ++ self.softmax_order_bar_0.arrive_and_wait() # g1 → g0, final even tile ++ else: ++ self.softmax_order_bar_0.arrive() # g1 → g0 + else: +- self.softmax_order_bar_1.arrive() # g0 → g1 ++ if is_order_chain_last_tile: ++ self.softmax_order_bar_1.arrive_and_wait() # g0 → g1, final odd tile ++ else: ++ self.softmax_order_bar_1.arrive() # g0 → g1 + # cute.nvgpu.cfence() + + mma_s_consumer_state.advance() +@@ -3547,6 +3825,24 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + + return mma_o_consumer_state + ++ @cute.jit ++ def dcp_split_has_valid_key( ++ self, ++ common_params: SimpleNamespace, ++ flat_q_row: cutlass.Int32, ++ ) -> cutlass.Boolean: ++ """Whether this row has a visible key in the current local K split.""" ++ first_local_key = common_params.k_index * self.mma_qk_tiler[1] ++ mask_threshold = self.num_heads * ( ++ self.cp_world * first_local_key ++ + common_params.cp_rank ++ - common_params.causal_seq_len ++ + self.seq_len_q ++ ) ++ return cute.elem_less(first_local_key, common_params.K) and not cute.elem_less( ++ flat_q_row, mask_threshold ++ ) ++ + @cute.jit + def epilogue( + self, +@@ -3600,14 +3896,33 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + cute.copy(tmem_load_tiled_copy, tTR_tAcc, tTR_rAcc) + + # apply output scale and normalize by row_sum +- for i in cutlass.range( +- cute.size(tTR_rAcc), vectorize=True, unroll_full=True +- ): +- tTR_rAcc[i] = ( +- tTR_rAcc[i] +- * epilogue_params.output_scale +- * cute.arch.rcp_approx(row_sum) ++ row_has_valid_key = True ++ if cutlass.const_expr(self.enable_dcp): ++ flat_q_row = ( ++ common_params.blk_coord[1] * self.mma_qk_tiler[0] + tTR_cO[0][0] ++ ) ++ row_has_valid_key = self.dcp_split_has_valid_key( ++ common_params, flat_q_row ++ ) ++ if cutlass.const_expr(self.enable_dcp): ++ output_normalizer = ( ++ epilogue_params.output_scale * cute.arch.rcp_approx(row_sum) ++ if row_has_valid_key ++ else self.acc_dtype(0.0) + ) ++ for i in cutlass.range( ++ cute.size(tTR_rAcc), vectorize=True, unroll_full=True ++ ): ++ tTR_rAcc[i] = tTR_rAcc[i] * output_normalizer ++ else: ++ for i in cutlass.range( ++ cute.size(tTR_rAcc), vectorize=True, unroll_full=True ++ ): ++ tTR_rAcc[i] = ( ++ tTR_rAcc[i] ++ * epilogue_params.output_scale ++ * cute.arch.rcp_approx(row_sum) ++ ) + + # store o to global memory + tR2G_rO_src = None +@@ -3688,6 +4003,8 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + cute.math.log2(row_sum, fastmath=True) + + epilogue_params.softmax_scale_log2 * row_max + ) ++ if cutlass.const_expr(self.enable_dcp): ++ lse = lse if row_has_valid_key else -self.lse_dtype.inf + # When writing directly to the user-facing mLSE (single-tile, + # no split-KV merge), convert from log2 base to natural log. + # When writing the per-split intermediate (mAccLSE branch), keep +@@ -3905,25 +4222,6 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + + return tile_sched_params, grid + +- @staticmethod +- def compute_fold_sq_ratio(num_heads: int, seq_len_q: int, m_tile: int) -> int: +- """Derive the seq_len_q-into-heads fold factor F. +- +- Returns the largest integer F such that: +- - F divides seq_len_q evenly +- - num_heads * F ≤ m_tile +- - 1 ≤ F ≤ seq_len_q +- +- F=1 means no folding (i.e. ``fold_sq`` should be False at the caller). +- """ +- if num_heads >= m_tile: +- return 1 +- max_fold = min(seq_len_q, m_tile // num_heads) +- for f in range(max_fold, 0, -1): +- if seq_len_q % f == 0: +- return f +- return 1 +- + @staticmethod + def get_workspace_size( + H: int, +@@ -3953,12 +4251,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + """ + if split_kv == 1: + return 0 +- # Decode packs heads into a physical 128-wide MMA-M tile. For H < 128, +- # split-KV partials can still touch the padded head lanes before +- # reduction, so size the workspace for max(H, 128). Mirrors the same +- # padding applied in initialize_workspace(). See #3235. +- workspace_heads = max(H, 128) +- return B * workspace_heads * S * split_kv * (D + 1) * acc_dtype.width // 8 ++ # The first workspace mode is the physical MMA-M row extent. Direct ++ # callers may still pass a logical H < 128, so preserve the M128 floor ++ # used by initialize_workspace(). The flat internal path passes 128. ++ workspace_rows = max(H, 128) ++ return B * workspace_rows * S * split_kv * (D + 1) * acc_dtype.width // 8 + + @cute.jit + def initialize_workspace( +@@ -3994,29 +4291,29 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + """ + acc_o, acc_lse = None, None + if cutlass.const_expr(workspace is not None): +- # Pad head dim to the physical 128-wide MMA-M tile. Without this, +- # H<128 split-KV partials write past the workspace. See #3235. +- workspace_H = cutlass.max(H, cutlass.Int32(128)) ++ # The first workspace mode is the physical MMA-M row extent. ++ # Preserve the M128 floor for direct callers that pass H < 128. ++ workspace_rows = cutlass.max(H, cutlass.Int32(128)) + align = 256 // self.q_dtype.width + acc_o_layout = cute.make_layout( +- (workspace_H, split_kv, D, S, B), ++ (workspace_rows, split_kv, D, S, B), + stride=( + cute.assume(split_kv * D, align), + cute.assume(D, align), + 1, +- cute.assume(split_kv * workspace_H * D, align), +- cute.assume(workspace_H * split_kv * S * D, align), ++ cute.assume(split_kv * workspace_rows * D, align), ++ cute.assume(workspace_rows * split_kv * S * D, align), + ), + ) + acc_o_iter = cute.recast_ptr(workspace.iterator, dtype=acc_dtype) + acc_o = cute.make_tensor(acc_o_iter, acc_o_layout) + acc_lse_layout = cute.make_layout( +- (workspace_H, split_kv, S, B), ++ (workspace_rows, split_kv, S, B), + stride=( + split_kv, + 1, +- workspace_H * split_kv, +- workspace_H * split_kv * S, ++ workspace_rows * split_kv, ++ workspace_rows * split_kv * S, + ), + ) + acc_lse_iter = cute.recast_ptr( +@@ -4098,11 +4395,11 @@ class BlackwellMultiHeadLatentAttentionForwardFP8: + return False + if is_var_split_kv and not is_var_seq: + return False +- if mma_qk_tiler_mn[0] < H: ++ if H <= 0 or mma_qk_tiler_mn[0] < H: + return False +- # When H < M tile, fold up to F tokens of S into H (M_eff = H*F ≤ M_tile). +- # F is auto-picked as the largest divisor of S with H*F ≤ M_tile. +- # F=1 always works, so any (H ≤ M_tile, S ≥ 1) is implementable. ++ # Query-token/head rows are flattened across 128-row M tiles. A ++ # partial final tile is supported, so any H <= M_tile and S >= 1 is ++ # valid. + if S < 1: + return False + if K <= 0: +diff --git a/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py b/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py +index 722bcfdf..0acd6919 100644 +--- a/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py ++++ b/flashinfer/cute_dsl/attention/monolithic/mla_helpers.py +@@ -302,3 +302,34 @@ MAX_SPLITS = 256 + + def ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b ++ ++ ++def compute_q_tile_layout( ++ num_heads: int, seq_len_q: int, m_tile: int = 128 ++) -> tuple[int, int, int]: ++ """Return ``(total_rows, num_tiles, tail_rows)`` for flat query packing. ++ ++ Query-token and head modes are one affine row space ordered as ++ ``flat_row = q_token * num_heads + q_head``. Consecutive M tiles may ++ therefore cross token boundaries; only the final tile can be partial. ++ ``tail_rows`` is always in ``[1, m_tile]`` and equals ``m_tile`` when the ++ flattened row count exactly fills the final tile. ++ ++ This host-side helper is shared by launch/workspace selection and both ++ kernel variants so their split-KV geometry cannot drift apart. ++ """ ++ if num_heads <= 0: ++ raise ValueError(f"num_heads must be positive, got {num_heads}") ++ if seq_len_q <= 0: ++ raise ValueError(f"seq_len_q must be positive, got {seq_len_q}") ++ if m_tile <= 0: ++ raise ValueError(f"m_tile must be positive, got {m_tile}") ++ if num_heads > m_tile: ++ raise ValueError( ++ f"num_heads ({num_heads}) must not exceed the MMA M tile ({m_tile})" ++ ) ++ ++ total_rows = num_heads * seq_len_q ++ num_tiles = ceil_div(total_rows, m_tile) ++ tail_rows = total_rows - (num_tiles - 1) * m_tile ++ return total_rows, num_tiles, tail_rows +diff --git a/flashinfer/mla/_core.py b/flashinfer/mla/_core.py +index bf422ec6..84b2c11f 100644 +--- a/flashinfer/mla/_core.py ++++ b/flashinfer/mla/_core.py +@@ -1983,13 +1983,66 @@ def _round_to_seq_len_bucket(x: int) -> int: + return 1 << (x - 1).bit_length() + + ++def _resolve_cute_dsl_workspace_sizer( ++ cute_dsl_impl: str, ++ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], ++ enable_dcp: bool = False, ++): ++ """Resolve the selected CuTeDSL implementation and its workspace policy.""" ++ from ..cute_dsl.attention.mla_dispatch import _resolve_impl ++ ++ resolved_impl = _resolve_impl( ++ requested=cute_dsl_impl, ++ kwargs={"sinks": sinks, "enable_dcp": enable_dcp}, ++ ) ++ if resolved_impl == "monolithic": ++ from ..cute_dsl.attention.monolithic.mla_decode import ( ++ _get_split_kv_and_workspace_size, ++ ) ++ else: ++ from ..cute_dsl.attention.wrappers.batch_mla import ( ++ _get_split_kv_and_workspace_size, ++ ) ++ return _get_split_kv_and_workspace_size, resolved_impl ++ ++ ++def _get_cute_dsl_workspace_sizer( ++ cute_dsl_impl: str, ++ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], ++ enable_dcp: bool = False, ++): ++ """Return the workspace policy owned by the selected CuTeDSL implementation.""" ++ return _resolve_cute_dsl_workspace_sizer(cute_dsl_impl, sinks, enable_dcp)[0] ++ ++ ++def _call_cute_dsl_workspace_sizer( ++ workspace_sizer, ++ resolved_impl: str, ++ batch_size: int, ++ q_len: int, ++ num_heads: int, ++ kv_lora_rank: int, ++ max_active_blocks: int, ++ max_seq_len: int, ++): ++ """Call an implementation's workspace policy with its supported arguments.""" ++ args = (batch_size, q_len, num_heads, kv_lora_rank, max_active_blocks) ++ if resolved_impl == "monolithic": ++ return workspace_sizer(*args, max_seq_len) ++ return workspace_sizer(*args) ++ ++ + def _cute_dsl_max_supported_batch( + workspace_bytes: int, + q_len: int, + num_heads: int, + kv_lora_rank: int, + max_active_blocks: int, ++ max_seq_len: int, + candidate_max: int, ++ cute_dsl_impl: str, ++ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], ++ enable_dcp: bool = False, + ) -> int: + """Largest batch the caller's workspace can support for cute-dsl MLA decode. + +@@ -1997,15 +2050,22 @@ def _cute_dsl_max_supported_batch( + split-K state. Binary-search for the largest ``B <= candidate_max`` whose + ``get_workspace_size(...)`` fits in ``workspace_bytes``. + """ +- from ..cute_dsl.attention.wrappers.batch_mla import ( +- _get_split_kv_and_workspace_size, ++ workspace_sizer, resolved_impl = _resolve_cute_dsl_workspace_sizer( ++ cute_dsl_impl, sinks, enable_dcp + ) + + lo, hi = 1, max(1, candidate_max) + while lo < hi: + mid = (lo + hi + 1) // 2 +- _, ws = _get_split_kv_and_workspace_size( +- mid, q_len, num_heads, kv_lora_rank, max_active_blocks ++ _, ws = _call_cute_dsl_workspace_sizer( ++ workspace_sizer, ++ resolved_impl, ++ mid, ++ q_len, ++ num_heads, ++ kv_lora_rank, ++ max_active_blocks, ++ max_seq_len, + ) + if ws <= workspace_bytes: + lo = mid +@@ -2020,7 +2080,11 @@ def _compute_mla_decode_buckets( + q_len: int, + num_heads: int, + kv_lora_rank: int, ++ max_seq_len: int, + device: torch.device, ++ cute_dsl_impl: str, ++ sinks: Optional[Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor]], ++ enable_dcp: bool = False, + ) -> Tuple[int, ...]: + """Compute the autotune bucket list from kernel/workspace limits only. + +@@ -2048,13 +2112,112 @@ def _compute_mla_decode_buckets( + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + max_active_blocks=get_num_sm(device), ++ max_seq_len=max_seq_len, + candidate_max=_TRTLLM_GEN_MLA_MAX_BATCH, ++ cute_dsl_impl=cute_dsl_impl, ++ sinks=sinks, ++ enable_dcp=enable_dcp, + ) + cap = max(cap, cute_dsl_cap) + + return get_hybrid_num_tokens_buckets(max(1, cap)) + + ++def _validate_mla_dcp_args( ++ *, ++ query: torch.Tensor, ++ backend: str, ++ sinks: Optional[List[torch.Tensor]], ++ cum_seq_lens_q: Optional[torch.Tensor], ++ max_q_len: Optional[int], ++ return_lse: bool, ++ enable_dcp: bool, ++ cp_world: int, ++ cp_rank: int, ++ causal_seqlens_kv_global: Optional[torch.Tensor], ++) -> str: ++ """Validate the public DCP contract and return the effective backend.""" ++ if not isinstance(enable_dcp, bool): ++ raise TypeError(f"enable_dcp must be a bool, got {type(enable_dcp).__name__}") ++ if not isinstance(cp_world, int) or isinstance(cp_world, bool) or cp_world <= 0: ++ raise ValueError(f"cp_world must be a positive integer, got {cp_world!r}") ++ if not isinstance(cp_rank, int) or isinstance(cp_rank, bool): ++ raise TypeError(f"cp_rank must be an integer, got {type(cp_rank).__name__}") ++ ++ if not enable_dcp: ++ nondefault = [] ++ if cp_world != 1: ++ nondefault.append(f"cp_world={cp_world}") ++ if cp_rank != 0: ++ nondefault.append(f"cp_rank={cp_rank}") ++ if causal_seqlens_kv_global is not None: ++ nondefault.append("causal_seqlens_kv_global") ++ if nondefault: ++ raise ValueError( ++ "DCP arguments require enable_dcp=True; got " + ", ".join(nondefault) ++ ) ++ return backend ++ ++ if query.ndim != 4: ++ raise ValueError( ++ "DCP requires a dense query with shape " ++ "[batch_size, q_len_per_request, num_heads, head_dim_qk]" ++ ) ++ if not 0 <= cp_rank < cp_world: ++ raise ValueError( ++ f"cp_rank must satisfy 0 <= cp_rank < cp_world, got " ++ f"cp_rank={cp_rank}, cp_world={cp_world}" ++ ) ++ if backend not in ("auto", "cute-dsl"): ++ raise ValueError( ++ f"enable_dcp=True is only supported by backend='cute-dsl', got " ++ f"backend={backend!r}" ++ ) ++ if not return_lse: ++ raise ValueError( ++ "enable_dcp=True requires return_lse=True so rank-local " ++ "attention states can be merged" ++ ) ++ if sinks is not None: ++ raise ValueError( ++ "DCP cannot be combined with sinks: DCP requires monolithic " ++ "CuTeDSL MLA, while sinks require the modular implementation" ++ ) ++ if cum_seq_lens_q is not None or max_q_len is not None: ++ raise ValueError("DCP does not support cum_seq_lens_q / max_q_len") ++ if causal_seqlens_kv_global is None: ++ raise ValueError( ++ "causal_seqlens_kv_global is required when enable_dcp=True" ++ ) ++ if not isinstance(causal_seqlens_kv_global, torch.Tensor): ++ raise TypeError( ++ "causal_seqlens_kv_global must be a torch.Tensor, got " ++ f"{type(causal_seqlens_kv_global).__name__}" ++ ) ++ if causal_seqlens_kv_global.dtype != torch.int32: ++ raise ValueError( ++ "causal_seqlens_kv_global must have dtype torch.int32, got " ++ f"{causal_seqlens_kv_global.dtype}" ++ ) ++ if not causal_seqlens_kv_global.is_cuda: ++ raise ValueError("causal_seqlens_kv_global must be a CUDA tensor") ++ if causal_seqlens_kv_global.device != query.device: ++ raise ValueError( ++ "causal_seqlens_kv_global must be on the query device " ++ f"{query.device}, got {causal_seqlens_kv_global.device}" ++ ) ++ if tuple(causal_seqlens_kv_global.shape) != (query.shape[0],): ++ raise ValueError( ++ "causal_seqlens_kv_global must have shape " ++ f"({query.shape[0]},), got {tuple(causal_seqlens_kv_global.shape)}" ++ ) ++ if not causal_seqlens_kv_global.is_contiguous(): ++ raise ValueError("causal_seqlens_kv_global must be contiguous") ++ ++ # DCP masking exists only in the monolithic CuTeDSL implementation. ++ return "cute-dsl" ++ ++ + def _cute_dsl_incompatibility_reason( + query: torch.Tensor, + out_dtype: torch.dtype, +@@ -2068,9 +2231,9 @@ def _cute_dsl_incompatibility_reason( + kv_lora_rank: int, + page_size: int, + is_var_seq: bool, +- return_lse: bool, +- lse: Optional[torch.Tensor], + cute_dsl_impl: str = "auto", ++ enable_dcp: bool = False, ++ cp_world: int = 1, + ) -> Optional[str]: + """Return None if cute-dsl can handle this call, else a human-readable reason. + +@@ -2123,8 +2286,14 @@ def _cute_dsl_incompatibility_reason( + try: + from ..cute_dsl.attention.mla_dispatch import _resolve_impl + +- resolved_impl = _resolve_impl(requested=cute_dsl_impl, kwargs={"sinks": sinks}) +- except (ValueError, ImportError) as e: ++ resolved_impl = _resolve_impl( ++ requested=cute_dsl_impl, ++ kwargs={ ++ "sinks": sinks, ++ "enable_dcp": enable_dcp, ++ }, ++ ) ++ except (TypeError, ValueError, ImportError) as e: + return f"cute-dsl backend (MLA decode kernel): {e}" + + try: +@@ -2133,7 +2302,7 @@ def _cute_dsl_incompatibility_reason( + else: + from ..cute_dsl.attention.wrappers.batch_mla import _check_can_implement + +- _check_can_implement( ++ check_kwargs = dict( + torch_dtype=query.dtype, + torch_out_dtype=out_dtype, + page_size=page_size, +@@ -2145,7 +2314,10 @@ def _cute_dsl_incompatibility_reason( + is_var_seq=is_var_seq, + is_var_split_kv=False, + ) +- except (ValueError, ImportError) as e: ++ if resolved_impl == "monolithic": ++ check_kwargs.update(enable_dcp=enable_dcp, cp_world=cp_world) ++ _check_can_implement(**check_kwargs) ++ except (TypeError, ValueError, ImportError) as e: + return f"cute-dsl backend (MLA decode kernel) cannot implement this configuration: {e}" + return None + +@@ -2155,8 +2327,11 @@ def _mla_decode_tuning_config( + buckets: tuple[int, ...], + num_pages: int, + profile_seq_len: int, ++ enable_dcp: bool = False, ++ cp_world: int = 1, ++ cp_rank: int = 0, + ) -> TuningConfig: +- """One TuningConfig (and one pair of initializer closures) per key. ++ """One TuningConfig and stable initializer set per key. + + Memoized because ``AutoTuner._find_nearest_profile`` lru-caches on + ``(shapes, tuning_config)``: a fresh config per dispatcher call shares +@@ -2164,12 +2339,13 @@ def _mla_decode_tuning_config( + ``tensor_initializers``) but never compares equal (closures compare by + identity), so would result in a leak. + +- The DynamicTensorSpec sweeps batch dim across all four ``inputs`` tensors +- (query, block_tables, seq_lens, out). ``block_tables`` is initialized via +- ``random_(0, num_pages)`` which wraps mod kv_cache size — safe for autotune +- profiling because MLA decode reads kv_cache and never writes it, so aliased +- page reads give correct timing measurements. ``seq_lens`` is filled +- homogeneously with ``profile_seq_len``. ++ The DynamicTensorSpec sweeps batch dim across ``query``, ``block_tables``, ++ ``seq_lens``, ``out``, and, for DCP, ``causal_seqlens_kv_global``. ++ ``block_tables`` is initialized via ``random_(0, num_pages)`` which wraps ++ mod kv_cache size — safe for autotune profiling because MLA decode reads ++ kv_cache and never writes it, so aliased page reads give correct timing ++ measurements. ``seq_lens`` is filled with ``profile_seq_len``; the ++ synthetic DCP global bound preserves that exact rank-local length. + """ + + def init_block_tables(shapes, dtype, device): +@@ -2182,14 +2358,32 @@ def _mla_decode_tuning_config( + tensor.fill_(profile_seq_len) + return tensor + ++ def init_causal_seqlens_kv_global(shapes, dtype, device): ++ tensor = torch.empty(shapes, dtype=dtype, device=device) ++ tensor.fill_(profile_seq_len * cp_world + cp_rank) ++ return tensor ++ ++ input_idx = (0, 1, 2, 3, 4) if enable_dcp else (0, 1, 2, 3) ++ tensor_initializers = ( ++ ( ++ None, ++ init_block_tables, ++ init_seq_lens, ++ None, ++ init_causal_seqlens_kv_global, ++ ) ++ if enable_dcp ++ else (None, init_block_tables, init_seq_lens, None) ++ ) ++ + return TuningConfig( + dynamic_tensor_specs=( + DynamicTensorSpec( +- input_idx=(0, 1, 2, 3), +- dim_idx=(0, 0, 0, 0), ++ input_idx=input_idx, ++ dim_idx=(0,) * len(input_idx), + gen_tuning_buckets=buckets, + map_to_tuning_buckets=make_bucket_mapper(buckets, round_map=False), +- tensor_initializers=(None, init_block_tables, init_seq_lens, None), ++ tensor_initializers=tensor_initializers, + ), + ), + use_cuda_graph=True, +@@ -2207,6 +2401,13 @@ def _build_mla_decode_tuning_config( + kv_lora_rank: int, + max_seq_len: int, + device: torch.device, ++ cute_dsl_impl: str = "auto", ++ sinks: Optional[ ++ Union[List[torch.Tensor], Tuple[torch.Tensor, ...], torch.Tensor] ++ ] = None, ++ enable_dcp: bool = False, ++ cp_world: int = 1, ++ cp_rank: int = 0, + ) -> TuningConfig: + """Reduce call args to the memoization key of ``_mla_decode_tuning_config``. + +@@ -2229,10 +2430,21 @@ def _build_mla_decode_tuning_config( + q_len, + num_heads, + kv_lora_rank, ++ max_seq_len, + device, ++ cute_dsl_impl, ++ sinks, ++ enable_dcp, + ) + +- return _mla_decode_tuning_config(buckets, num_pages, profile_seq_len) ++ return _mla_decode_tuning_config( ++ buckets, ++ num_pages, ++ profile_seq_len, ++ enable_dcp, ++ cp_world, ++ cp_rank, ++ ) + + + class TrtllmGenMlaDecodeRunner(TunableRunner): +@@ -2458,6 +2670,9 @@ class CuteDslMlaDecodeRunner(TunableRunner): + return_lse: bool, + sinks: Optional[torch.Tensor], + cute_dsl_impl: str, ++ enable_dcp: bool = False, ++ cp_world: int = 1, ++ cp_rank: int = 0, + ): + from ..cute_dsl.attention import cute_dsl_mla_decode + +@@ -2482,6 +2697,13 @@ class CuteDslMlaDecodeRunner(TunableRunner): + self.return_lse = return_lse + self.sinks = sinks + self.cute_dsl_impl = cute_dsl_impl ++ self.enable_dcp = enable_dcp ++ self.cp_world = cp_world ++ self.cp_rank = cp_rank ++ self._profile_lse: Optional[torch.Tensor] = None ++ self._workspace_sizer, self._resolved_cute_dsl_impl = ( ++ _resolve_cute_dsl_workspace_sizer(cute_dsl_impl, sinks, enable_dcp) ++ ) + + def __hash__(self): + # See TrtllmGenMlaDecodeRunner.__hash__ — tactic-determining state is +@@ -2495,15 +2717,19 @@ class CuteDslMlaDecodeRunner(TunableRunner): + # If the caller's workspace can't fit batch=B for this profile, opt + # out so the autotuner skips us (no JIT cost) and trtllm-gen wins by + # default for that bucket. +- from ..cute_dsl.attention.wrappers.batch_mla import ( +- _get_split_kv_and_workspace_size, +- ) + from ..cute_dsl.utils import get_num_sm + + q = inputs[0] + B, q_len, num_heads, _ = q.shape +- _, ws = _get_split_kv_and_workspace_size( +- B, q_len, num_heads, self.kv_lora_rank, get_num_sm(q.device) ++ _, ws = _call_cute_dsl_workspace_sizer( ++ self._workspace_sizer, ++ self._resolved_cute_dsl_impl, ++ B, ++ q_len, ++ num_heads, ++ self.kv_lora_rank, ++ get_num_sm(q.device), ++ self.max_seq_len, + ) + workspace_bytes = ( + self.workspace_buffer.numel() * self.workspace_buffer.element_size() +@@ -2513,14 +2739,24 @@ class CuteDslMlaDecodeRunner(TunableRunner): + return [-1] + + def get_cache_key_extras(self, inputs): +- q, _, _, out = inputs ++ q, _, _, out = inputs[:4] + # Cute-dsl rejects sparse/skip-softmax/tensor-scales upstream, so + # those are omitted from extras as constants for this runner. + # ``sinks`` and ``cute_dsl_impl`` are included because they flip the +- # impl (modular vs monolithic) inside ``cute_dsl_mla_decode``. ++ # impl (modular vs monolithic) inside ``cute_dsl_mla_decode``. The ++ # monolithic K-tile count and workspace capacity keep cache hits from ++ # bypassing a different split-workspace validity decision. + sinks_key = ( + None if self.sinks is None else (tuple(self.sinks.shape), self.sinks.dtype) + ) ++ workspace_bytes = ( ++ self.workspace_buffer.numel() * self.workspace_buffer.element_size() ++ ) ++ seq_len_workspace_key = ( ++ (self.max_seq_len + 127) // 128 ++ if self._resolved_cute_dsl_impl == "monolithic" ++ else _round_to_seq_len_bucket(self.max_seq_len) ++ ) + return ( + q.dtype, + self.kv_cache.dtype, +@@ -2529,12 +2765,15 @@ class CuteDslMlaDecodeRunner(TunableRunner): + self.kv_lora_rank, + self.qk_rope_head_dim, + self.page_size, +- _round_to_seq_len_bucket(self.max_seq_len), ++ seq_len_workspace_key, ++ workspace_bytes, + self.is_var_seq, + self.uses_shared_paged_kv_idx, + self.enable_pdl, + sinks_key, + self.cute_dsl_impl, ++ getattr(self, "enable_dcp", False), ++ getattr(self, "cp_world", 1), + ) + + def forward( +@@ -2544,7 +2783,31 @@ class CuteDslMlaDecodeRunner(TunableRunner): + do_preparation: bool = False, + **kwargs, + ): +- query, block_tables, seq_lens, out = inputs ++ query, block_tables, seq_lens, out = inputs[:4] ++ causal_seqlens_kv_global = inputs[4] if self.enable_dcp else None ++ ++ # LSE is not a tuning input because it does not influence tactic ++ # selection. When a synthetic batch differs from the caller batch, ++ # provide matching temporary storage while retaining the caller's LSE ++ # for the final invocation. ++ lse = self.lse ++ if self.return_lse: ++ expected_numel = query.shape[0] * query.shape[1] * query.shape[2] ++ if lse is None or lse.numel() != expected_numel: ++ expected_shape = ( ++ query.shape[0] * query.shape[1], ++ query.shape[2], ++ ) ++ if ( ++ self._profile_lse is None ++ or tuple(self._profile_lse.shape) != expected_shape ++ ): ++ self._profile_lse = torch.empty( ++ expected_shape, ++ dtype=torch.float32, ++ device=query.device, ++ ) ++ lse = self._profile_lse + return self._run( + query=query, + kv_cache=self.kv_cache, +@@ -2560,10 +2823,14 @@ class CuteDslMlaDecodeRunner(TunableRunner): + out_dtype=self.out_dtype, + is_var_seq=self.is_var_seq, + enable_pdl=self.enable_pdl, +- lse=self.lse, ++ lse=lse, + return_lse=self.return_lse, + sinks=self.sinks, + cute_dsl_impl=self.cute_dsl_impl, ++ enable_dcp=self.enable_dcp, ++ cp_world=self.cp_world, ++ cp_rank=self.cp_rank, ++ causal_seqlens_kv_global=causal_seqlens_kv_global, + ) + + +@@ -2595,6 +2862,10 @@ def trtllm_batch_decode_with_kv_cache_mla( + cum_seq_lens_q: Optional[torch.Tensor] = None, + max_q_len: Optional[int] = None, + multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None, ++ enable_dcp: bool = False, ++ cp_world: int = 1, ++ cp_rank: int = 0, ++ causal_seqlens_kv_global: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + r"""Decode MLA with TRTLLM-GEN, CuteDSL, XQA, or SM120/SM121 sparse kernels. + +@@ -2639,12 +2910,15 @@ def trtllm_batch_decode_with_kv_cache_mla( + With ``backend="trtllm-gen"``, the final dimension may use its native + width and does not need padding to a multiple of ``128 / page_size``. + seq_lens : Optional[torch.Tensor] +- Per-request KV sequence lengths for dense and TRTLLM-GEN paths. For ++ Per-request physical KV sequence lengths for dense and TRTLLM-GEN ++ paths. With DCP these are rank-local lengths and continue to control ++ paging, memory bounds, and split-KV. For + SM120/SM121 sparse v32/GLM, pass ``[batch_size, q_len_per_request]`` or + flattened ``[batch_size * q_len_per_request]`` active top-k lengths; if + ``None``, every column in ``block_tables`` is active. + max_seq_len : int +- Maximum KV sequence length used for dense/TRTLLM-GEN scheduling. ++ Maximum physical KV sequence length used for dense/TRTLLM-GEN ++ scheduling. With DCP this is the maximum rank-local length. + Ignored by the SM120/SM121 sparse v32/GLM backend. + sparse_mla_top_k : int + Enables sparse MLA when greater than zero. On SM100/SM103 this selects +@@ -2756,6 +3030,18 @@ def trtllm_batch_decode_with_kv_cache_mla( + for each concurrently executing CUDA stream or graph. Autotune profiling + uses runner-owned internal storage; the caller buffer is used only for the + final request. ++ enable_dcp : bool = False ++ Statically enable cyclic decode context parallelism in the monolithic ++ CuTeDSL MLA kernel. DCP returns a rank-local output/LSE state, so ++ ``return_lse=True`` is required and the caller must merge rank states. ++ cp_world : int = 1 ++ Compile-time context-parallel world size. Rank ``r`` stores global KV ++ positions whose token index modulo ``cp_world`` equals ``r``. ++ cp_rank : int = 0 ++ Runtime-uniform context-parallel rank. ++ causal_seqlens_kv_global : Optional[torch.Tensor] = None ++ Contiguous CUDA int32 tensor ``[batch_size]`` containing the global ++ exclusive causal bound for the newest query token. Required with DCP. + + Note + ---- +@@ -2779,11 +3065,12 @@ def trtllm_batch_decode_with_kv_cache_mla( + Autotune + -------- + On SM100/SM103 dense MLA, calling under ``flashinfer.autotune(True)`` with +- ``backend="auto"`` profiles both ``trtllm-gen`` and ``cute-dsl`` across a +- bucketed batch sweep up to each runner's kernel/workspace cap and caches the +- winning runner per shape signature. Subsequent calls under +- ``autotune(False)`` dispatch to the cached choice; any batch outside the +- tuned range falls back to a default runner with a one-time warning. ++ ``backend="auto"`` profiles both ``trtllm-gen`` and ``cute-dsl`` when DCP ++ is disabled. DCP profiles only its required monolithic CuTeDSL runner. ++ Both modes use a bucketed batch sweep up to each runner's kernel/workspace ++ cap and cache the winning tactic per shape signature. Subsequent calls under ++ ``autotune(False)`` use the cached choice; any batch outside the tuned range ++ falls back to the runner's default tactic with a one-time warning. + + The autotune bucket range and cache key do **not** depend on + ``kv_cache.shape[0]`` (the number of pages in the pool), so reallocating the +@@ -2805,6 +3092,19 @@ def trtllm_batch_decode_with_kv_cache_mla( + if max_q_len is not None and cum_seq_lens_q is None: + raise ValueError("max_q_len is only supported when cum_seq_lens_q is provided") + ++ backend = _validate_mla_dcp_args( ++ query=query, ++ backend=backend, ++ sinks=sinks, ++ cum_seq_lens_q=cum_seq_lens_q, ++ max_q_len=max_q_len, ++ return_lse=return_lse, ++ enable_dcp=enable_dcp, ++ cp_world=cp_world, ++ cp_rank=cp_rank, ++ causal_seqlens_kv_global=causal_seqlens_kv_global, ++ ) ++ + if backend == "auto": + cc = get_compute_capability(query.device) + if cc[0] == 12 and sparse_mla_top_k > 0: +@@ -3116,9 +3416,9 @@ def trtllm_batch_decode_with_kv_cache_mla( + kv_lora_rank, + page_size, + is_var_seq, +- return_lse, +- lse, + cute_dsl_impl, ++ enable_dcp, ++ cp_world, + ) + if backend == "cute-dsl": + if cute_dsl_reason is not None: +@@ -3203,6 +3503,9 @@ def trtllm_batch_decode_with_kv_cache_mla( + return_lse=return_lse, + sinks=cute_dsl_sinks, + cute_dsl_impl=cute_dsl_impl, ++ enable_dcp=enable_dcp, ++ cp_world=cp_world, ++ cp_rank=cp_rank, + ) + ) + +@@ -3217,8 +3520,17 @@ def trtllm_batch_decode_with_kv_cache_mla( + kv_lora_rank=kv_lora_rank, + max_seq_len=max_seq_len, + device=query.device, ++ cute_dsl_impl=cute_dsl_impl, ++ sinks=sinks, ++ enable_dcp=enable_dcp, ++ cp_world=cp_world, ++ cp_rank=cp_rank, + ) + inputs = [query, block_tables, seq_lens, out] ++ if enable_dcp: ++ # The global causal bound varies with batch and must be synthesized ++ # alongside the other batch-shaped tensors during autotuning. ++ inputs.append(causal_seqlens_kv_global) + runner, tactic = AutoTuner.get().choose_one( + "trtllm_batch_decode_mla", + runners, +diff --git a/flashinfer/trace/templates/attention.py b/flashinfer/trace/templates/attention.py +index 610afb82..7b32bb88 100644 +--- a/flashinfer/trace/templates/attention.py ++++ b/flashinfer/trace/templates/attention.py +@@ -2692,6 +2692,11 @@ trtllm_batch_decode_mla_sparse_trace = TraceTemplate( + + + def trtllm_batch_decode_mla_trace_dispatch(**kwargs): ++ if kwargs.get("enable_dcp", False): ++ raise NotImplementedError( ++ "fi_trace does not yet represent cyclic DCP KV ownership or " ++ "cross-rank LSE merging for MLA decode" ++ ) + sparse_mla_top_k = int(kwargs.get("sparse_mla_top_k", 0) or 0) + if sparse_mla_top_k > 0: + return trtllm_batch_decode_mla_sparse_trace +diff --git a/tests/attention/test_cute_dsl_mla_dcp.py b/tests/attention/test_cute_dsl_mla_dcp.py +new file mode 100644 +index 00000000..ee1bfabe +--- /dev/null ++++ b/tests/attention/test_cute_dsl_mla_dcp.py +@@ -0,0 +1,1126 @@ ++# Copyright (c) 2026 by FlashInfer team. ++# ++# Licensed under the Apache License, Version 2.0 (the "License"); ++# you may not use this file except in compliance with the License. ++# You may obtain a copy of the License at ++# ++# http://www.apache.org/licenses/LICENSE-2.0 ++# ++# Unless required by applicable law or agreed to in writing, software ++# distributed under the License is distributed on an "AS IS" BASIS, ++# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. ++# See the License for the specific language governing permissions and ++# limitations under the License. ++ ++"""Correctness tests for static DCP support in monolithic CuTe DSL MLA.""" ++ ++import math ++from bisect import bisect_left ++from unittest import mock ++ ++import pytest ++import torch ++ ++from flashinfer.cute_dsl import is_cute_dsl_available ++from flashinfer.utils import is_sm100a_supported, is_sm110a_supported ++ ++ ++_LATENT_DIM = 512 ++_ROPE_DIM = 64 ++_QK_DIM = _LATENT_DIM + _ROPE_DIM ++_PAGE_SIZE = 64 ++ ++ ++def _skip_if_unsupported() -> None: ++ device = torch.device("cuda") ++ if not (is_sm100a_supported(device) or is_sm110a_supported(device)): ++ pytest.skip("Requires SM100-SM110 (tcgen05)") ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ ++ ++def _ceil_div(numerator: int, denominator: int) -> int: ++ return -(-numerator // denominator) ++ ++ ++def _local_causal_bound( ++ global_bound_newest: int, ++ q_len: int, ++ q_idx: int, ++ cp_world: int, ++ cp_rank: int, ++) -> int: ++ return _ceil_div( ++ global_bound_newest - cp_rank - (q_len - 1) + q_idx, ++ cp_world, ++ ) ++ ++ ++def _flat_dcp_score_is_valid( ++ flat_q_row: int, ++ num_heads: int, ++ local_key: int, ++ q_len: int, ++ global_bound_newest: int, ++ cp_world: int, ++ cp_rank: int, ++) -> bool: ++ return flat_q_row >= num_heads * ( ++ cp_world * local_key ++ + cp_rank ++ - global_bound_newest ++ + q_len ++ ) ++ ++ ++def _local_length(global_length: int, cp_world: int, cp_rank: int) -> int: ++ return max(_ceil_div(global_length - cp_rank, cp_world), 0) ++ ++ ++def _make_inputs( ++ *, ++ global_length: int, ++ q_len: int, ++ num_heads: int, ++ dtype: torch.dtype, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ++ """Create one batch of quantized-once Q and global-coordinate KV.""" ++ return _make_batched_inputs( ++ global_lengths=(global_length,), ++ q_len=q_len, ++ num_heads=num_heads, ++ dtype=dtype, ++ ) ++ ++ ++def _make_batched_inputs( ++ *, ++ global_lengths: tuple[int, ...], ++ q_len: int, ++ num_heads: int, ++ dtype: torch.dtype, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ++ """Create a heterogeneous batch with one padded global-coordinate KV pool.""" ++ device = torch.device("cuda") ++ storage_dtype = torch.float16 if dtype == torch.float8_e4m3fn else dtype ++ batch_size = len(global_lengths) ++ query = ( ++ torch.randn( ++ batch_size, ++ q_len, ++ num_heads, ++ _QK_DIM, ++ device=device, ++ dtype=storage_dtype, ++ ) ++ * 0.1 ++ ).to(dtype) ++ global_kv = ( ++ torch.randn( ++ batch_size, ++ max(global_lengths), ++ _QK_DIM, ++ device=device, ++ dtype=storage_dtype, ++ ) ++ * 0.1 ++ ).to(dtype) ++ global_lens = torch.tensor( ++ global_lengths, dtype=torch.int32, device=device ++ ) ++ return query, global_kv, global_lens ++ ++ ++def _pack_cyclic_rank_pages( ++ global_kv: torch.Tensor, ++ global_lens: torch.Tensor, ++ cp_world: int, ++ cp_rank: int, ++) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: ++ """Pack ``g % cp_world == cp_rank`` tokens into a contiguous paged cache.""" ++ batch_size, _, d_qk = global_kv.shape ++ global_lens_host = global_lens.tolist() ++ local_lens_host = [ ++ _local_length(global_len, cp_world, cp_rank) ++ for global_len in global_lens_host ++ ] ++ pages_per_batch = [ ++ max(1, _ceil_div(local_len, _PAGE_SIZE)) ++ for local_len in local_lens_host ++ ] ++ max_pages = max(pages_per_batch) ++ total_pages = sum(pages_per_batch) ++ ++ local_cache = torch.zeros( ++ total_pages, ++ _PAGE_SIZE, ++ d_qk, ++ dtype=global_kv.dtype, ++ device=global_kv.device, ++ ) ++ block_tables = torch.zeros( ++ batch_size, ++ max_pages, ++ dtype=torch.int32, ++ device=global_kv.device, ++ ) ++ ++ next_page = 0 ++ for batch_idx, (global_len, local_len, num_pages) in enumerate( ++ zip( ++ global_lens_host, ++ local_lens_host, ++ pages_per_batch, ++ strict=True, ++ ) ++ ): ++ page_ids = torch.arange( ++ next_page, ++ next_page + num_pages, ++ dtype=torch.int32, ++ device=global_kv.device, ++ ) ++ block_tables[batch_idx, :num_pages] = page_ids ++ if local_len: ++ local_tokens = global_kv[ ++ batch_idx, cp_rank:global_len:cp_world ++ ] ++ local_cache[next_page : next_page + num_pages].view( ++ -1, d_qk ++ )[:local_len].copy_(local_tokens) ++ next_page += num_pages ++ ++ local_lens = torch.tensor( ++ local_lens_host, dtype=torch.int32, device=global_kv.device ++ ) ++ return ( ++ local_cache, ++ block_tables, ++ local_lens, ++ max(1, max(local_lens_host)), ++ ) ++ ++ ++def _reference_attention( ++ query: torch.Tensor, ++ global_kv: torch.Tensor, ++ global_lens: torch.Tensor, ++ *, ++ cp_world: int = 1, ++ cp_rank: int = 0, ++) -> tuple[torch.Tensor, torch.Tensor]: ++ """Reference one cyclic rank; world=1 is the full-context reference.""" ++ batch_size, q_len, num_heads, _ = query.shape ++ out = torch.zeros( ++ batch_size, ++ q_len, ++ num_heads, ++ _LATENT_DIM, ++ dtype=torch.float32, ++ device=query.device, ++ ) ++ lse = torch.full( ++ (batch_size, q_len, num_heads), ++ -torch.inf, ++ dtype=torch.float32, ++ device=query.device, ++ ) ++ softmax_scale = 1.0 / math.sqrt(_LATENT_DIM) ++ ++ for batch_idx, global_len in enumerate(global_lens.tolist()): ++ keys = global_kv[ ++ batch_idx, cp_rank:global_len:cp_world ++ ].float() ++ global_positions = range(cp_rank, global_len, cp_world) ++ for q_idx in range(q_len): ++ global_bound = global_len - (q_len - 1) + q_idx ++ # Count directly in global coordinates so this reference remains ++ # independent of the ceiling-divided kernel formula without ++ # synchronizing on a CUDA predicate. ++ visible_count = bisect_left(global_positions, global_bound) ++ if visible_count == 0: ++ continue ++ visible_keys = keys[:visible_count] ++ scores = torch.einsum( ++ "hd,kd->hk", ++ query[batch_idx, q_idx].float(), ++ visible_keys, ++ ) ++ scores *= softmax_scale ++ lse[batch_idx, q_idx] = torch.logsumexp(scores, dim=-1) ++ probabilities = torch.softmax(scores, dim=-1) ++ out[batch_idx, q_idx] = torch.einsum( ++ "hk,kd->hd", ++ probabilities, ++ visible_keys[:, :_LATENT_DIM], ++ ) ++ return out, lse ++ ++ ++def _merge_rank_outputs_natural_log( ++ rank_outputs: list[torch.Tensor], ++ rank_lses: list[torch.Tensor], ++) -> tuple[torch.Tensor, torch.Tensor]: ++ """Stable LSE-weighted merge for the kernel's natural-log public LSE.""" ++ outputs = torch.stack([out.float() for out in rank_outputs], dim=0) ++ lses = torch.stack([lse.float() for lse in rank_lses], dim=0) ++ max_lse = lses.max(dim=0).values ++ has_keys = torch.isfinite(max_lse) ++ safe_max = torch.where(has_keys, max_lse, torch.zeros_like(max_lse)) ++ weights = torch.where( ++ torch.isfinite(lses), ++ torch.exp(lses - safe_max.unsqueeze(0)), ++ torch.zeros_like(lses), ++ ) ++ weight_sum = weights.sum(dim=0) ++ safe_sum = torch.where(has_keys, weight_sum, torch.ones_like(weight_sum)) ++ merged_out = ( ++ outputs * weights.unsqueeze(-1) ++ ).sum(dim=0) / safe_sum.unsqueeze(-1) ++ merged_out = torch.where( ++ has_keys.unsqueeze(-1), merged_out, torch.zeros_like(merged_out) ++ ) ++ merged_lse = torch.where( ++ has_keys, ++ safe_max + torch.log(safe_sum), ++ torch.full_like(max_lse, -torch.inf), ++ ) ++ return merged_out, merged_lse ++ ++ ++def _prepare_rank_call( ++ query: torch.Tensor, ++ global_kv: torch.Tensor, ++ global_lens: torch.Tensor, ++ *, ++ cp_world: int, ++ cp_rank: int, ++ is_var_seq: bool = False, ++ enable_pdl: bool = False, ++) -> tuple[dict, int]: ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _get_split_kv_and_workspace_size, ++ ) ++ from flashinfer.cute_dsl.utils import get_num_sm ++ ++ kv_cache, block_tables, local_lens, max_local_len = ( ++ _pack_cyclic_rank_pages( ++ global_kv, global_lens, cp_world, cp_rank ++ ) ++ ) ++ batch_size, q_len, num_heads, _ = query.shape ++ split_kv, workspace_size = _get_split_kv_and_workspace_size( ++ batch_size, ++ q_len, ++ num_heads, ++ _LATENT_DIM, ++ get_num_sm(query.device), ++ max_local_len, ++ ) ++ workspace = torch.empty( ++ max(workspace_size, 1), dtype=torch.int8, device=query.device ++ ) ++ return ( ++ { ++ "query": query, ++ "kv_cache": kv_cache, ++ "workspace_buffer": workspace, ++ "kv_lora_rank": _LATENT_DIM, ++ "qk_rope_head_dim": _ROPE_DIM, ++ "block_tables": block_tables, ++ "seq_lens": local_lens, ++ "max_seq_len": max_local_len, ++ "softmax_scale": 1.0 / math.sqrt(_LATENT_DIM), ++ "is_var_seq": is_var_seq, ++ "enable_pdl": enable_pdl, ++ }, ++ split_kv, ++ ) ++ ++ ++def _launch_rank( ++ query: torch.Tensor, ++ global_kv: torch.Tensor, ++ global_lens: torch.Tensor, ++ *, ++ cp_world: int, ++ cp_rank: int, ++ enable_dcp: bool, ++ causal_lens: torch.Tensor | None = None, ++ is_var_seq: bool = False, ++ enable_pdl: bool = False, ++) -> tuple[torch.Tensor, torch.Tensor, int]: ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ cute_dsl_mla_decode, ++ ) ++ ++ call_args, split_kv = _prepare_rank_call( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=cp_world, ++ cp_rank=cp_rank, ++ is_var_seq=is_var_seq, ++ enable_pdl=enable_pdl, ++ ) ++ kwargs = {} ++ if enable_dcp: ++ kwargs = { ++ "enable_dcp": True, ++ "cp_world": cp_world, ++ "cp_rank": cp_rank, ++ "causal_seqlens_kv_global": ( ++ global_lens if causal_lens is None else causal_lens ++ ), ++ } ++ out, lse = cute_dsl_mla_decode( ++ return_lse=True, ++ **call_args, ++ **kwargs, ++ ) ++ return out, lse, split_kv ++ ++ ++def _assert_close_to_reference( ++ out: torch.Tensor, ++ lse: torch.Tensor, ++ ref_out: torch.Tensor, ++ ref_lse: torch.Tensor, ++ dtype: torch.dtype, ++) -> None: ++ if dtype == torch.float8_e4m3fn: ++ out_atol, out_rtol = 0.1, 0.1 ++ lse_atol, lse_rtol = 0.2, 0.1 ++ else: ++ out_atol = out_rtol = lse_atol = lse_rtol = 1e-2 ++ torch.testing.assert_close( ++ out.float(), ref_out, atol=out_atol, rtol=out_rtol ++ ) ++ torch.testing.assert_close( ++ lse.float(), ref_lse, atol=lse_atol, rtol=lse_rtol ++ ) ++ ++ ++def _assert_dcp_rank_merge_matches_reference( ++ query: torch.Tensor, ++ global_kv: torch.Tensor, ++ global_lens: torch.Tensor, ++ *, ++ cp_world: int, ++ dtype: torch.dtype, ++ is_var_seq: bool = False, ++) -> list[int]: ++ """Check every rank-local state and their final natural-log merge.""" ++ rank_outputs = [] ++ rank_lses = [] ++ split_kvs = [] ++ for cp_rank in range(cp_world): ++ out, lse, split_kv = _launch_rank( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=cp_world, ++ cp_rank=cp_rank, ++ enable_dcp=True, ++ is_var_seq=is_var_seq, ++ ) ++ ref_rank_out, ref_rank_lse = _reference_attention( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=cp_world, ++ cp_rank=cp_rank, ++ ) ++ _assert_close_to_reference( ++ out, lse, ref_rank_out, ref_rank_lse, dtype ++ ) ++ rank_outputs.append(out) ++ rank_lses.append(lse) ++ split_kvs.append(split_kv) ++ ++ merged_out, merged_lse = _merge_rank_outputs_natural_log( ++ rank_outputs, rank_lses ++ ) ++ ref_out, ref_lse = _reference_attention( ++ query, global_kv, global_lens ++ ) ++ _assert_close_to_reference( ++ merged_out, merged_lse, ref_out, ref_lse, dtype ++ ) ++ return split_kvs ++ ++ ++def test_dcp_flat_mask_matches_ceiling_divided_bound(): ++ """The division-free flattened predicate must match global coordinates.""" ++ for num_heads in (6, 12, 24, 48, 64, 96, 128): ++ for q_len in (1, 2, 4, 8): ++ for cp_world in (1, 2, 4): ++ for cp_rank in range(cp_world): ++ for global_len in (q_len, q_len + 1, 127, 128, 129): ++ local_len = _local_length( ++ global_len, cp_world, cp_rank ++ ) ++ for q_idx in range(q_len): ++ for local_key in range(local_len): ++ expected = ( ++ local_key * cp_world + cp_rank ++ < global_len - (q_len - 1) + q_idx ++ ) ++ for head in (0, num_heads - 1): ++ flat_q_row = q_idx * num_heads + head ++ assert ( ++ _flat_dcp_score_is_valid( ++ flat_q_row, ++ num_heads, ++ local_key, ++ q_len, ++ global_len, ++ cp_world, ++ cp_rank, ++ ) ++ == expected ++ ) ++ ++ ++def test_dcp_per_query_tile_dense_boundary_is_conservative(): ++ """A tile marked dense must be valid for every real row and local key.""" ++ k_tile = 128 ++ for num_heads in (6, 12, 24, 48, 64, 96, 128): ++ for q_len in (2, 4, 8, 32): ++ num_q_tiles = _ceil_div(q_len * num_heads, 128) ++ for cp_world in (1, 2, 4): ++ for cp_rank in range(cp_world): ++ global_len = max(q_len, 277) ++ local_len = _local_length( ++ global_len, cp_world, cp_rank ++ ) ++ global_positions = range( ++ cp_rank, global_len, cp_world ++ ) ++ for q_tile_idx in range(num_q_tiles): ++ first_flat_row = q_tile_idx * 128 ++ first_q = first_flat_row // num_heads ++ earliest_bound = bisect_left( ++ global_positions, ++ global_len - (q_len - 1) + first_q, ++ ) ++ effective_bound = min( ++ max(earliest_bound, 0), local_len ++ ) ++ first_mask_k_tile = effective_bound // k_tile ++ for k_tile_idx in range(first_mask_k_tile): ++ local_begin = k_tile_idx * k_tile ++ local_end = min(local_begin + k_tile, local_len) ++ for flat_q_row in range( ++ first_flat_row, ++ min( ++ first_flat_row + 128, ++ q_len * num_heads, ++ ), ++ ): ++ for local_key in range( ++ local_begin, local_end ++ ): ++ assert _flat_dcp_score_is_valid( ++ flat_q_row, ++ num_heads, ++ local_key, ++ q_len, ++ global_len, ++ cp_world, ++ cp_rank, ++ ) ++ ++def test_cute_dsl_mla_dcp_rejects_incomplete_static_contract(): ++ """Reject DCP calls that cannot produce mergeable rank-local states.""" ++ _skip_if_unsupported() ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ cute_dsl_mla_decode, ++ ) ++ ++ torch.manual_seed(40) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=128, ++ q_len=4, ++ num_heads=96, ++ dtype=torch.bfloat16, ++ ) ++ call_args, _ = _prepare_rank_call( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=0, ++ ) ++ ++ ragged_call_args = {**call_args, "query": query[:, 0]} ++ with pytest.raises(ValueError, match="query must have shape"): ++ cute_dsl_mla_decode(**ragged_call_args) ++ with pytest.raises(ValueError, match="requires return_lse=True"): ++ cute_dsl_mla_decode( ++ **call_args, ++ enable_dcp=True, ++ cp_world=2, ++ cp_rank=0, ++ causal_seqlens_kv_global=global_lens, ++ ) ++ with pytest.raises( ++ ValueError, match="causal_seqlens_kv_global is required" ++ ): ++ cute_dsl_mla_decode( ++ **call_args, ++ enable_dcp=True, ++ cp_world=2, ++ cp_rank=0, ++ return_lse=True, ++ ) ++ with pytest.raises(TypeError, match="must be a torch.Tensor"): ++ cute_dsl_mla_decode( ++ **call_args, ++ enable_dcp=True, ++ cp_world=2, ++ cp_rank=0, ++ causal_seqlens_kv_global=[128], ++ return_lse=True, ++ ) ++ with pytest.raises(ValueError, match="0 <= cp_rank < cp_world"): ++ cute_dsl_mla_decode( ++ **call_args, ++ enable_dcp=True, ++ cp_world=2, ++ cp_rank=2, ++ causal_seqlens_kv_global=global_lens, ++ return_lse=True, ++ ) ++ with pytest.raises(ValueError, match="require enable_dcp=True"): ++ cute_dsl_mla_decode(**call_args, cp_world=2) ++ ++ ++def test_cute_dsl_mla_dcp_dispatch_is_strictly_monolithic(): ++ """DCP must never silently select a backend with non-DCP mask semantics.""" ++ from flashinfer.cute_dsl.attention.mla_dispatch import _resolve_impl ++ from flashinfer.trace.templates.attention import ( ++ trtllm_batch_decode_mla_trace_dispatch, ++ ) ++ ++ assert ( ++ _resolve_impl( ++ requested="auto", ++ kwargs={"enable_dcp": True, "cp_world": 2, "cp_rank": 0}, ++ ) ++ == "monolithic" ++ ) ++ with pytest.raises(ValueError, match="only supported by the monolithic"): ++ _resolve_impl( ++ requested="modular", ++ kwargs={"enable_dcp": True, "cp_world": 2, "cp_rank": 0}, ++ ) ++ with pytest.raises(ValueError, match="cannot be combined with 'sinks'"): ++ _resolve_impl( ++ requested="auto", ++ kwargs={ ++ "enable_dcp": True, ++ "cp_world": 2, ++ "cp_rank": 0, ++ "sinks": torch.empty(1), ++ }, ++ ) ++ with pytest.raises(NotImplementedError, match="does not yet represent"): ++ trtllm_batch_decode_mla_trace_dispatch(enable_dcp=True) ++ ++ ++def test_cute_dsl_mla_dcp_world1_matches_disabled(): ++ """World-one DCP must preserve the existing monolithic MLA result.""" ++ _skip_if_unsupported() ++ torch.manual_seed(41) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=128, ++ q_len=4, ++ num_heads=96, ++ dtype=torch.bfloat16, ++ ) ++ baseline_out, baseline_lse, baseline_split = _launch_rank( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=1, ++ cp_rank=0, ++ enable_dcp=False, ++ ) ++ dcp_out, dcp_lse, dcp_split = _launch_rank( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=1, ++ cp_rank=0, ++ enable_dcp=True, ++ ) ++ assert baseline_split == dcp_split == 1 ++ torch.testing.assert_close(dcp_out, baseline_out, atol=0, rtol=0) ++ torch.testing.assert_close(dcp_lse, baseline_lse, atol=0, rtol=0) ++ ++ ++@pytest.mark.parametrize( ++ "dtype", ++ [torch.bfloat16, torch.float8_e4m3fn], ++ ids=["bf16", "fp8"], ++) ++def test_cute_dsl_mla_dcp_q1_h128_rank_merge(dtype): ++ """Cover the distinct Q1/H128 schedule for both kernel families.""" ++ _skip_if_unsupported() ++ torch.manual_seed(49) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=129, ++ q_len=1, ++ num_heads=128, ++ dtype=dtype, ++ ) ++ split_kvs = _assert_dcp_rank_merge_matches_reference( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ dtype=dtype, ++ ) ++ assert split_kvs == [1, 1] ++ ++ ++@pytest.mark.parametrize("num_heads", [12, 24, 48]) ++def test_cute_dsl_mla_dcp_packed_head_counts(num_heads): ++ """Exercise DCP across the supported packed query-tile geometries.""" ++ _skip_if_unsupported() ++ torch.manual_seed(50 + num_heads) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=130, ++ q_len=8, ++ num_heads=num_heads, ++ dtype=torch.bfloat16, ++ ) ++ split_kvs = _assert_dcp_rank_merge_matches_reference( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ dtype=torch.bfloat16, ++ ) ++ assert split_kvs == [1, 1] ++ ++ ++def test_cute_dsl_mla_dcp_public_api_autotune_profiles_causal_tensor(): ++ """The public runner must batch-sweep DCP metadata and return caller B1.""" ++ _skip_if_unsupported() ++ from flashinfer import autotune ++ from flashinfer.autotuner import AutoTuner ++ from flashinfer.mla._core import ( ++ trtllm_batch_decode_with_kv_cache_mla, ++ ) ++ ++ torch.manual_seed(47) ++ query, global_kv, global_lens = _make_inputs( ++ # Two local pages satisfy the common public dispatcher's aligned ++ # block-table contract for page_size=64. ++ global_length=256, ++ q_len=4, ++ num_heads=96, ++ dtype=torch.bfloat16, ++ ) ++ call_args, split_kv = _prepare_rank_call( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=0, ++ ) ++ assert split_kv == 1 ++ ++ public_args = { ++ "query": query, ++ "kv_cache": call_args["kv_cache"], ++ "workspace_buffer": call_args["workspace_buffer"], ++ "qk_nope_head_dim": 128, ++ "kv_lora_rank": _LATENT_DIM, ++ "qk_rope_head_dim": _ROPE_DIM, ++ "block_tables": call_args["block_tables"], ++ "seq_lens": call_args["seq_lens"], ++ "max_seq_len": call_args["max_seq_len"], ++ "bmm1_scale": call_args["softmax_scale"], ++ "bmm2_scale": 1.0, ++ "backend": "auto", ++ "cute_dsl_impl": "monolithic", ++ "is_var_seq": False, ++ "enable_pdl": False, ++ "return_lse": True, ++ "enable_dcp": True, ++ "cp_world": 2, ++ "cp_rank": 0, ++ } ++ with pytest.raises(TypeError, match="must be a torch.Tensor"): ++ trtllm_batch_decode_with_kv_cache_mla( ++ **public_args, ++ causal_seqlens_kv_global=[256], ++ ) ++ ++ AutoTuner.get().clear_cache() ++ try: ++ with mock.patch( ++ "flashinfer.mla._core._compute_mla_decode_buckets", ++ return_value=(2,), ++ ), autotune(tune_mode=True): ++ out, lse = trtllm_batch_decode_with_kv_cache_mla( ++ **public_args, ++ causal_seqlens_kv_global=global_lens, ++ ) ++ finally: ++ AutoTuner.get().clear_cache() ++ ++ ref_out, ref_lse = _reference_attention( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=0, ++ ) ++ _assert_close_to_reference( ++ out, ++ lse.view_as(ref_lse), ++ ref_out, ++ ref_lse, ++ torch.bfloat16, ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "dtype", ++ [torch.bfloat16, torch.float8_e4m3fn], ++ ids=["bf16", "fp8"], ++) ++def test_cute_dsl_mla_dcp_rank_mask_and_merge(dtype): ++ """Cover the rank-sensitive G128/Q4/H96 boundary and rank merge.""" ++ _skip_if_unsupported() ++ torch.manual_seed(42) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=128, ++ q_len=4, ++ num_heads=96, ++ dtype=dtype, ++ ) ++ split_kvs = _assert_dcp_rank_merge_matches_reference( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ dtype=dtype, ++ ) ++ assert split_kvs == [1, 1] ++ ++ # In particular, query zero on rank one has local bound 62, not 63. ++ assert ( ++ _local_causal_bound(128, 4, 0, cp_world=2, cp_rank=1) ++ == 62 ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "dtype", ++ [torch.float16, torch.float8_e4m3fn], ++ ids=["fp16", "fp8"], ++) ++def test_cute_dsl_mla_dcp_world4_variable_batch_h6(dtype): ++ """Cover W4, heterogeneous local tails, FP16, and packed H6 query rows.""" ++ _skip_if_unsupported() ++ torch.manual_seed(46) ++ query, global_kv, global_lens = _make_batched_inputs( ++ global_lengths=(65, 130, 259), ++ q_len=8, ++ num_heads=6, ++ dtype=dtype, ++ ) ++ split_kvs = _assert_dcp_rank_merge_matches_reference( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=4, ++ dtype=dtype, ++ is_var_seq=True, ++ ) ++ assert split_kvs == [1] * 4 ++ ++ ++@pytest.mark.parametrize( ++ "cp_world,global_length,dtype", ++ [ ++ pytest.param(4, 4 * 1024 + 3, torch.bfloat16, id="w4-k4k-bf16"), ++ pytest.param( ++ 4, ++ 4 * 1024 + 3, ++ torch.float8_e4m3fn, ++ id="w4-k4k-fp8", ++ ), ++ pytest.param(8, 8 * 1024 + 5, torch.bfloat16, id="w8-k8k-bf16"), ++ pytest.param( ++ 8, ++ 8 * 1024 + 5, ++ torch.float8_e4m3fn, ++ id="w8-k8k-fp8", ++ ), ++ pytest.param(16, 16 * 1024 + 9, torch.bfloat16, id="w16-k16k-bf16"), ++ pytest.param( ++ 16, ++ 16 * 1024 + 9, ++ torch.float8_e4m3fn, ++ id="w16-k16k-fp8", ++ ), ++ ], ++) ++def test_cute_dsl_mla_dcp_rank_scale_merge( ++ cp_world, global_length, dtype ++): ++ """Cover W4/W8/W16 four-query speculative decode and uneven split-KV.""" ++ _skip_if_unsupported() ++ torch.manual_seed(51 + cp_world) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=global_length, ++ q_len=4, ++ num_heads=96, ++ dtype=dtype, ++ ) ++ ++ base_local_length, tail_ranks = divmod(global_length, cp_world) ++ for cp_rank in range(cp_world): ++ expected_local_length = base_local_length + ( ++ 1 if cp_rank < tail_ranks else 0 ++ ) ++ assert ( ++ _local_length(global_length, cp_world, cp_rank) ++ == expected_local_length ++ ) ++ ++ split_kvs = _assert_dcp_rank_merge_matches_reference( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=cp_world, ++ dtype=dtype, ++ ) ++ assert all(split_kv > 1 for split_kv in split_kvs) ++ ++ ++def test_cute_dsl_mla_dcp_cuda_graph_reads_updated_causal_bound(): ++ """A captured launch must read causal bounds mutated in place on replay.""" ++ _skip_if_unsupported() ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ cute_dsl_mla_decode, ++ ) ++ ++ torch.manual_seed(45) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=128, ++ q_len=4, ++ num_heads=96, ++ dtype=torch.bfloat16, ++ ) ++ call_args, split_kv = _prepare_rank_call( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=1, ++ ) ++ assert split_kv == 1 ++ out = torch.empty( ++ 1, ++ 4, ++ 96, ++ _LATENT_DIM, ++ dtype=torch.bfloat16, ++ device=query.device, ++ ) ++ lse = torch.empty( ++ 1, 4, 96, dtype=torch.float32, device=query.device ++ ) ++ dcp_args = { ++ "enable_dcp": True, ++ "cp_world": 2, ++ "cp_rank": 1, ++ "causal_seqlens_kv_global": global_lens, ++ "return_lse": True, ++ "out": out, ++ "lse": lse, ++ } ++ ++ # Compile and initialize all persistent buffers before capture. ++ returned_out, returned_lse = cute_dsl_mla_decode( ++ **call_args, **dcp_args ++ ) ++ assert returned_out is out ++ assert returned_lse is lse ++ torch.cuda.synchronize() ++ ++ graph = torch.cuda.CUDAGraph() ++ with torch.cuda.graph(graph): ++ cute_dsl_mla_decode(**call_args, **dcp_args) ++ lse_at_128 = lse.clone() ++ ++ # Keep the physical local cache and seq_lens fixed at the G=128 shard. ++ # The G=126 replay must mask its now-extra global keys 126 and 127. ++ global_lens.fill_(126) ++ graph.replay() ++ torch.cuda.synchronize() ++ assert not torch.equal(lse, lse_at_128) ++ ref_out, ref_lse = _reference_attention( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=1, ++ ) ++ _assert_close_to_reference( ++ out, lse, ref_out, ref_lse, torch.bfloat16 ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "dtype", ++ [torch.bfloat16, torch.float8_e4m3fn], ++ ids=["bf16", "fp8"], ++) ++def test_cute_dsl_mla_dcp_empty_rank_row(dtype): ++ """A rank with no visible key must contribute O=0 and LSE=-inf.""" ++ _skip_if_unsupported() ++ torch.manual_seed(43) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=4, ++ q_len=4, ++ num_heads=96, ++ dtype=dtype, ++ ) ++ out, lse, split_kv = _launch_rank( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=1, ++ enable_dcp=True, ++ ) ++ assert split_kv == 1 ++ assert torch.equal(out[:, 0], torch.zeros_like(out[:, 0])) ++ assert torch.isneginf(lse[:, 0]).all() ++ ref_out, ref_lse = _reference_attention( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=1, ++ ) ++ _assert_close_to_reference( ++ out, lse, ref_out, ref_lse, dtype ++ ) ++ ++ # Also cover a physically empty cyclic shard and the cross-rank merge for ++ # early rows that have no visible key on any rank. The empty shard still ++ # owns one padding page, but seq_lens=0, so it may not load or reduce K. ++ query, global_kv, global_lens = _make_inputs( ++ global_length=1, ++ q_len=4, ++ num_heads=96, ++ dtype=dtype, ++ ) ++ split_kvs = _assert_dcp_rank_merge_matches_reference( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ dtype=dtype, ++ ) ++ assert split_kvs == [1, 1] ++ ++ ++@pytest.mark.parametrize( ++ "dtype", ++ [torch.bfloat16, torch.float8_e4m3fn], ++ ids=["bf16", "fp8"], ++) ++def test_cute_dsl_mla_dcp_split_kv_rank_merge(dtype): ++ """DCP must preserve natural-log LSE through standalone split reduction.""" ++ _skip_if_unsupported() ++ torch.manual_seed(44) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=4096, ++ q_len=4, ++ num_heads=96, ++ dtype=dtype, ++ ) ++ split_kvs = _assert_dcp_rank_merge_matches_reference( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ dtype=dtype, ++ ) ++ assert all(split_kv > 1 for split_kv in split_kvs) ++ ++ # Preserve the 2048-token physical shard and its split geometry, but make ++ # every rank-1 key causally invisible. This exercises the all-empty split ++ # reducer rather than the direct zero-tile output path above. ++ empty_causal_lens = torch.tensor( ++ [1], dtype=torch.int32, device=query.device ++ ) ++ empty_out, empty_lse, empty_split_kv = _launch_rank( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=1, ++ enable_dcp=True, ++ causal_lens=empty_causal_lens, ++ ) ++ assert empty_split_kv == split_kvs[1] > 1 ++ assert torch.equal(empty_out, torch.zeros_like(empty_out)) ++ assert torch.isneginf(empty_lse).all() ++ ref_empty_out, ref_empty_lse = _reference_attention( ++ query, ++ global_kv, ++ empty_causal_lens, ++ cp_world=2, ++ cp_rank=1, ++ ) ++ _assert_close_to_reference( ++ empty_out, ++ empty_lse, ++ ref_empty_out, ++ ref_empty_lse, ++ dtype, ++ ) ++ ++ ++def test_cute_dsl_mla_dcp_pdl_all_empty_split_reduction(): ++ """PDL must order the all-empty producer path before split reduction.""" ++ _skip_if_unsupported() ++ from flashinfer.utils import device_support_pdl ++ ++ if not device_support_pdl(torch.device("cuda")): ++ pytest.skip("Programmatic dependent launch is not supported") ++ ++ torch.manual_seed(48) ++ query, global_kv, global_lens = _make_inputs( ++ global_length=4096, ++ q_len=4, ++ num_heads=96, ++ dtype=torch.bfloat16, ++ ) ++ empty_causal_lens = torch.tensor( ++ [1], dtype=torch.int32, device=query.device ++ ) ++ out, lse, split_kv = _launch_rank( ++ query, ++ global_kv, ++ global_lens, ++ cp_world=2, ++ cp_rank=1, ++ enable_dcp=True, ++ causal_lens=empty_causal_lens, ++ enable_pdl=True, ++ ) ++ assert split_kv > 1 ++ assert torch.equal(out, torch.zeros_like(out)) ++ assert torch.isneginf(lse).all() +diff --git a/tests/attention/test_cute_dsl_mla_decode.py b/tests/attention/test_cute_dsl_mla_decode.py +index e8549a72..f52dee4f 100644 +--- a/tests/attention/test_cute_dsl_mla_decode.py ++++ b/tests/attention/test_cute_dsl_mla_decode.py +@@ -83,15 +83,19 @@ def torch_reference_mla( + + outputs = [] + lses = [] ++ # Copy the small metadata arrays once. Calling CUDA ``.item()`` for every ++ # request and page serializes hundreds of device synchronizations at B128. ++ cache_seqs_host = cache_seqs.tolist() ++ page_table_host = page_table.tolist() + for b in range(B): +- seq_len = cache_seqs[b].item() ++ seq_len = cache_seqs_host[b] + num_pages_needed = (seq_len + page_size - 1) // page_size + + # Gather KV for this batch via page table +- page_indices = page_table[b, :num_pages_needed] ++ page_indices = page_table_host[b][:num_pages_needed] + kv_indices = [] + for p in page_indices: +- start = p.item() * page_size ++ start = p * page_size + kv_indices.extend(range(start, start + page_size)) + kv_indices = kv_indices[:seq_len] + kv_indices_t = torch.tensor(kv_indices, device=q_nope.device) +@@ -254,28 +258,21 @@ def test_cute_dsl_mla_decode_fp16( + torch.testing.assert_close(lse, ref_lse, atol=1e-2, rtol=1e-2) + + +-# Exercises the spec-decoding (MTP) causal mask + fold_sq path: num_heads < 128 +-# forces the kernel to pack F = compute_fold_sq_ratio(H, q_len, 128) tokens of +-# q_len into the head dim so the 128-wide MMA-M tile is fully populated. +-# (H=128, q_len=any) → F=1 (no fold), (H=64, q_len=2) → F=2, (H=64, q_len=4) → F=2, +-# (H=32, q_len=4) → F=4, (H=32, q_len=2) → F=2. All paths share the same +-# kernel; the MTP causal mask is applied uniformly for q_len > 1. +-# Monolithic-only: the modular path doesn't implement fold_sq or the MTP mask. +-@pytest.mark.parametrize("batch_size", [1, 4]) +-@pytest.mark.parametrize("seq_len_k", [128, 1024]) +-@pytest.mark.parametrize("num_heads", [16, 32, 64]) +-@pytest.mark.parametrize("q_len", [2, 4]) +-@pytest.mark.parametrize("dtype", [torch.float16, torch.float8_e4m3fn]) +-def test_cute_dsl_mla_decode_fold_sq( +- batch_size, seq_len_k, num_heads, q_len, dtype, cute_dsl_impl ++def _run_padded_q_tile_case( ++ batch_size, ++ seq_len_k, ++ num_heads, ++ q_len, ++ dtype, ++ is_var_seq=False, ++ enable_pdl=None, ++ out_dtype=None, ++ query_token_stride=1, ++ via_public_api=False, + ): +- """Verify the MTP causal mask + fold_sq packing for H ≤ 128 and q_len > 1.""" +- if cute_dsl_impl != "monolithic": +- pytest.skip("fold_sq / MTP causal mask are monolithic-only features") ++ """Run and reference-check one monolithic packed-query configuration.""" + skip_if_unsupported() + +- from flashinfer.cute_dsl.attention import cute_dsl_mla_decode +- + torch.manual_seed(42) + device = torch.device("cuda") + +@@ -288,17 +285,19 @@ def test_cute_dsl_mla_decode_fold_sq( + + # torch.randn doesn't support fp8; for FP8 inputs create as fp16 then convert. + is_fp8 = dtype == torch.float8_e4m3fn ++ query_storage = torch.randn( ++ batch_size, ++ q_len * query_token_stride, ++ num_heads, ++ D_qk, ++ dtype=torch.float16 if is_fp8 else dtype, ++ device=device, ++ ) + if is_fp8: +- query = ( +- torch.randn( +- batch_size, q_len, num_heads, D_qk, dtype=torch.float16, device=device +- ) +- * 0.1 +- ).to(torch.float8_e4m3fn) +- else: +- query = torch.randn( +- batch_size, q_len, num_heads, D_qk, dtype=dtype, device=device +- ) ++ query_storage = (query_storage * 0.1).to(torch.float8_e4m3fn) ++ query = query_storage[:, ::query_token_stride] ++ if query_token_stride != 1: ++ assert query.stride(1) != num_heads * query.stride(2) + + num_pages_per_batch = (seq_len_k + page_size - 1) // page_size + total_pages = num_pages_per_batch * batch_size + 10 +@@ -319,25 +318,79 @@ def test_cute_dsl_mla_decode_fold_sq( + for p in range(num_pages_per_batch): + block_tables[b, p] = b * num_pages_per_batch + p + +- seq_lens = torch.full((batch_size,), seq_len_k, dtype=torch.int32, device=device) +- +- workspace_buffer = torch.empty(256 * 1024 * 1024, dtype=torch.int8, device=device) ++ if is_var_seq: ++ seq_lens = torch.tensor( ++ [max(page_size, seq_len_k - b * 37) for b in range(batch_size)], ++ dtype=torch.int32, ++ device=device, ++ ) ++ else: ++ seq_lens = torch.full( ++ (batch_size,), seq_len_k, dtype=torch.int32, device=device ++ ) + +- out = cute_dsl_mla_decode( +- query=query, +- kv_cache=kv_cache, +- workspace_buffer=workspace_buffer, +- kv_lora_rank=latent_dim, +- qk_rope_head_dim=rope_dim, +- block_tables=block_tables, +- seq_lens=seq_lens, +- max_seq_len=seq_len_k, +- softmax_scale=softmax_scale, +- output_scale=output_scale, +- is_var_seq=False, +- cute_dsl_impl=cute_dsl_impl, ++ workspace_factory = torch.zeros if via_public_api else torch.empty ++ workspace_buffer = workspace_factory( ++ 256 * 1024 * 1024, dtype=torch.int8, device=device + ) + ++ if via_public_api: ++ assert out_dtype is None ++ from flashinfer.mla import trtllm_batch_decode_with_kv_cache_mla ++ ++ lse_out = torch.empty( ++ batch_size, q_len, num_heads, dtype=torch.float32, device=device ++ ) ++ out, lse = trtllm_batch_decode_with_kv_cache_mla( ++ query=query, ++ kv_cache=kv_cache, ++ workspace_buffer=workspace_buffer, ++ qk_nope_head_dim=latent_dim, ++ kv_lora_rank=latent_dim, ++ qk_rope_head_dim=rope_dim, ++ block_tables=block_tables, ++ seq_lens=seq_lens, ++ max_seq_len=seq_len_k, ++ bmm1_scale=softmax_scale, ++ bmm2_scale=output_scale, ++ backend="cute-dsl", ++ is_var_seq=is_var_seq, ++ enable_pdl=enable_pdl, ++ lse=lse_out, ++ return_lse=True, ++ cute_dsl_impl="monolithic", ++ ) ++ else: ++ from flashinfer.cute_dsl.attention import cute_dsl_mla_decode ++ ++ out, lse = cute_dsl_mla_decode( ++ query=query, ++ kv_cache=kv_cache, ++ workspace_buffer=workspace_buffer, ++ kv_lora_rank=latent_dim, ++ qk_rope_head_dim=rope_dim, ++ block_tables=block_tables, ++ seq_lens=seq_lens, ++ max_seq_len=seq_len_k, ++ softmax_scale=softmax_scale, ++ output_scale=output_scale, ++ is_var_seq=is_var_seq, ++ enable_pdl=enable_pdl, ++ cute_dsl_impl="monolithic", ++ return_lse=True, ++ out_dtype=out_dtype, ++ ) ++ ++ if via_public_api: ++ expected_out_dtype = torch.bfloat16 ++ elif out_dtype is not None: ++ expected_out_dtype = out_dtype ++ elif is_fp8: ++ expected_out_dtype = torch.bfloat16 ++ else: ++ expected_out_dtype = dtype ++ assert out.dtype == expected_out_dtype ++ + # FP8 input → BF16 output (default), so do the reference in FP32 with wider tolerance. + if is_fp8: + kv_flat = kv_cache.reshape(-1, D_qk).to(torch.float32) +@@ -351,7 +404,7 @@ def test_cute_dsl_mla_decode_fold_sq( + c_rope_ref = kv_flat[:, latent_dim:] + + # Monolithic-only test — kernel always applies the MTP causal mask here. +- ref_out = torch_reference_mla( ++ ref_out, ref_lse = torch_reference_mla( + q_nope, + q_rope, + c_latent_ref, +@@ -362,6 +415,7 @@ def test_cute_dsl_mla_decode_fold_sq( + output_scale, + page_size, + apply_mtp_mask=True, ++ return_lse=True, + ) + + if is_fp8: +@@ -369,45 +423,518 @@ def test_cute_dsl_mla_decode_fold_sq( + torch.testing.assert_close( + out.to(torch.float32), ref_out.to(torch.float32), atol=0.1, rtol=0.1 + ) ++ torch.testing.assert_close(lse, ref_lse, atol=0.2, rtol=0.1) + else: +- ref_out_cast = ref_out.to(dtype) ++ ref_out_cast = ref_out.to(out.dtype) + torch.testing.assert_close(out, ref_out_cast, atol=1e-2, rtol=1e-2) ++ torch.testing.assert_close(lse, ref_lse, atol=1e-2, rtol=1e-2) ++ ++ ++# Exercises the spec-decoding (MTP) causal mask + flat query-row packing. ++# A cooperative 2-CTA tile owns 128 consecutive (token, head) rows and may ++# cross token boundaries; only the final flattened tile is safely padded. ++# All paths share the same kernel, and the MTP mask applies for q_len > 1. ++# Monolithic-only: the modular path doesn't implement packed query tiles or MTP. ++@pytest.mark.parametrize( ++ "num_heads,q_len,dtype", ++ [ ++ pytest.param(64, 3, torch.bfloat16, id="h64-sq3-bf16"), ++ pytest.param(24, 13, torch.float8_e4m3fn, id="h24-sq13-fp8"), ++ pytest.param(12, 11, torch.bfloat16, id="h12-sq11-bf16-tail4"), ++ ], ++) ++def test_cute_dsl_mla_decode_padded_q_tiles(num_heads, q_len, dtype): ++ """Cover cross-token rows and final tails in both kernel families.""" ++ _run_padded_q_tile_case(1, 1024, num_heads, q_len, dtype) ++ ++ ++@pytest.mark.parametrize( ++ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] ++) ++def test_cute_dsl_mla_decode_per_q_tile_mask_boundary(dtype): ++ """Use each M128 tile's first token when selecting the dense mask path.""" ++ skip_if_unsupported() ++ ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _get_split_kv_and_workspace_size, ++ cute_dsl_mla_decode, ++ ) ++ from flashinfer.cute_dsl.utils import get_num_sm ++ ++ batch_size, q_len, num_heads = 1, 32, 6 ++ seq_len_k, page_size = 144, 64 ++ latent_dim, rope_dim = 512, 64 ++ d_qk = latent_dim + rope_dim ++ device = torch.device("cuda") ++ ++ # H6/Sq32 has two M128 tiles. Token 21 straddles their boundary: heads ++ # 0-1 are in tile 0 and heads 2-5 are in tile 1. Tile 1 can treat keys ++ # 0:128 as dense from its first token (21), while tile 0 must still mask ++ # keys 113:128 for its earlier tokens. Using tile 0's last token for the ++ # coarse decision would expose those marked values incorrectly. ++ query = torch.zeros(batch_size, q_len, num_heads, d_qk, dtype=dtype, device=device) ++ num_pages = (seq_len_k + page_size - 1) // page_size ++ kv_cache = torch.zeros(num_pages, page_size, d_qk, dtype=dtype, device=device) ++ kv_cache.view(-1, d_qk)[113:128, 0] = 4.0 ++ block_tables = torch.arange(num_pages, dtype=torch.int32, device=device).view(1, -1) ++ seq_lens = torch.full((batch_size,), seq_len_k, dtype=torch.int32, device=device) ++ ++ split_kv, workspace_size = _get_split_kv_and_workspace_size( ++ batch_size, ++ q_len, ++ num_heads, ++ latent_dim, ++ get_num_sm(device), ++ seq_len_k, ++ ) ++ assert split_kv == 2 ++ workspace = torch.empty(workspace_size, dtype=torch.int8, device=device) ++ ++ out, lse = cute_dsl_mla_decode( ++ query=query, ++ kv_cache=kv_cache, ++ workspace_buffer=workspace, ++ kv_lora_rank=latent_dim, ++ qk_rope_head_dim=rope_dim, ++ block_tables=block_tables, ++ seq_lens=seq_lens, ++ max_seq_len=seq_len_k, ++ softmax_scale=1.0 / (latent_dim**0.5), ++ return_lse=True, ++ enable_pdl=False, ++ ) ++ ++ bounds = torch.arange(q_len, dtype=torch.float32, device=device) + ( ++ seq_len_k - q_len + 1 ++ ) ++ marked_visible = (bounds - 113).clamp(min=0, max=15) ++ expected_out = torch.zeros_like(out, dtype=torch.float32) ++ expected_out[..., 0] = (4.0 * marked_visible / bounds).view(1, q_len, 1) ++ expected_lse = torch.log(bounds).view(1, q_len, 1).expand_as(lse) ++ ++ torch.testing.assert_close(out.float(), expected_out, atol=5e-3, rtol=5e-3) ++ torch.testing.assert_close(lse, expected_lse, atol=2e-3, rtol=2e-3) ++ ++ ++@pytest.mark.parametrize("num_heads", [12, 24, 48, 96]) ++@pytest.mark.parametrize("q_len", [1, 2, 4, 8]) ++@pytest.mark.parametrize( ++ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] ++) ++def test_cute_dsl_mla_decode_packed_q_accuracy_matrix(num_heads, q_len, dtype): ++ """Qualify the supported packed-query, causal-mask, and auto-split matrix.""" ++ # Keep the batch sweep within one pytest item while all 128 requested ++ # Cartesian cells are reference checked. Split-KV and query packing are ++ # selected automatically by the public launcher. ++ for batch_size in (1, 4, 16, 128): ++ _run_padded_q_tile_case( ++ batch_size=batch_size, ++ seq_len_k=1024, ++ num_heads=num_heads, ++ q_len=q_len, ++ dtype=dtype, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_fp8_persistent_multi_work_boundaries(): ++ """Reuse the FP8 two-softmax pipelines across persistent work items.""" ++ # B128 exceeds one resident wave of 2CTA clusters on supported Blackwell ++ # devices. These cases exercise one-, even-, and odd-K-tile work items; ++ # the packed-Q matrix above additionally covers the eight-tile boundary. ++ for seq_len_k, enable_pdl in ((128, True), (256, False), (384, True)): ++ _run_padded_q_tile_case( ++ batch_size=128, ++ seq_len_k=seq_len_k, ++ num_heads=96, ++ q_len=1, ++ dtype=torch.float8_e4m3fn, ++ enable_pdl=enable_pdl, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_fp8_variable_seq_order_boundaries(): ++ """Balance both FP8 softmax groups for mixed nonpersistent K parity.""" ++ _run_padded_q_tile_case( ++ batch_size=10, ++ seq_len_k=385, ++ num_heads=96, ++ q_len=8, ++ dtype=torch.float8_e4m3fn, ++ is_var_seq=True, ++ enable_pdl=False, ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] ++) ++def test_cute_dsl_mla_decode_h96_max_split_reducer_capacity(dtype): ++ """Reference-check output and LSE at the static reducer's 32-split cap.""" ++ _run_padded_q_tile_case( ++ batch_size=1, ++ seq_len_k=8192, ++ num_heads=96, ++ q_len=1, ++ dtype=dtype, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_h96_sq8_nonempty_split_reducer(): ++ """Reference-check the normalized H96/Sq8 long-K split-reducer path.""" ++ _run_padded_q_tile_case( ++ batch_size=1, ++ seq_len_k=8192, ++ num_heads=96, ++ q_len=8, ++ dtype=torch.bfloat16, ++ ) ++ ++ ++@pytest.mark.parametrize( ++ "dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"] ++) ++def test_cute_dsl_mla_decode_h96_odd_split_reducer_pdl_off(dtype): ++ """Cover adaptive D4 with an odd 17-split prefix and PDL disabled.""" ++ _run_padded_q_tile_case( ++ batch_size=1, ++ seq_len_k=4097, ++ num_heads=96, ++ q_len=1, ++ dtype=dtype, ++ enable_pdl=False, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_variable_seq_d2_reducer(): ++ """Cover adaptive D2 when batches have distinct non-power split prefixes.""" ++ _run_padded_q_tile_case( ++ batch_size=2, ++ seq_len_k=4097, ++ num_heads=24, ++ q_len=1, ++ dtype=torch.bfloat16, ++ is_var_seq=True, ++ enable_pdl=False, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_padded_q_tile_direct_output(): ++ """Exercise H48/Sq3 full/tail padding with split_kv=1 (no reducer).""" ++ skip_if_unsupported() ++ ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _get_split_kv_and_workspace_size, ++ ) ++ from flashinfer.cute_dsl.utils import get_num_sm ++ ++ num_q_tiles = 2 ++ num_sm = get_num_sm(torch.device("cuda")) ++ batch_size = (num_sm + 2 * num_q_tiles - 1) // (2 * num_q_tiles) ++ split_kv, workspace_size = _get_split_kv_and_workspace_size( ++ batch_size, 3, 48, 512, num_sm ++ ) ++ assert split_kv == 1 ++ assert workspace_size == 0 ++ ++ _run_padded_q_tile_case( ++ batch_size=batch_size, ++ seq_len_k=128, ++ num_heads=48, ++ q_len=3, ++ dtype=torch.float16, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_padded_q_tile_variable_seq(): ++ """Cover nonpersistent variable-K scheduling with a padded H24/Sq6 tail.""" ++ _run_padded_q_tile_case( ++ batch_size=3, ++ seq_len_k=385, ++ num_heads=24, ++ q_len=6, ++ dtype=torch.float16, ++ is_var_seq=True, ++ enable_pdl=False, ++ ) + + +-def test_compute_fold_sq_ratio(): +- """Unit test the static helper used by both run() and the wrapper.""" ++def test_cute_dsl_mla_decode_padded_q_tile_fp8_output(): ++ """Qualify packed-query FP8 input and FP8 output together.""" ++ _run_padded_q_tile_case( ++ batch_size=1, ++ seq_len_k=128, ++ num_heads=64, ++ q_len=3, ++ dtype=torch.float8_e4m3fn, ++ out_dtype=torch.float8_e4m3fn, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_padded_q_tile_strided_query(): ++ """Cover the contiguous fallback for a token-strided query view.""" ++ _run_padded_q_tile_case( ++ batch_size=1, ++ seq_len_k=128, ++ num_heads=48, ++ q_len=3, ++ dtype=torch.float16, ++ query_token_stride=2, ++ ) ++ ++ ++@pytest.mark.parametrize("buffer_name", ["out", "lse"]) ++def test_cute_dsl_mla_decode_rejects_token_gapped_output(buffer_name): ++ """Reject output views that cannot represent flat rows across tokens.""" + if not is_cute_dsl_available(): + pytest.skip("CuTe DSL not available") +- from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp16 import ( +- BlackwellMultiHeadLatentAttentionForwardFP16 as FP16, ++ ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ cute_dsl_mla_decode, + ) +- from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp8 import ( +- BlackwellMultiHeadLatentAttentionForwardFP8 as FP8, ++ ++ batch_size, q_len, num_heads = 1, 2, 96 ++ latent_dim, rope_dim = 512, 64 ++ query = torch.empty( ++ batch_size, ++ q_len, ++ num_heads, ++ latent_dim + rope_dim, ++ dtype=torch.bfloat16, ++ ) ++ out = torch.empty(batch_size, q_len * 2, num_heads, latent_dim)[:, ::2] ++ lse = torch.empty(batch_size, q_len * 2, num_heads, dtype=torch.float32)[:, ::2] ++ kwargs = {buffer_name: out if buffer_name == "out" else lse} ++ ++ with pytest.raises(ValueError, match=rf"{buffer_name} must be contiguous"): ++ cute_dsl_mla_decode( ++ query=query, ++ kv_cache=torch.empty(1, 64, latent_dim + rope_dim, dtype=torch.bfloat16), ++ workspace_buffer=torch.empty(0, dtype=torch.int8), ++ kv_lora_rank=latent_dim, ++ qk_rope_head_dim=rope_dim, ++ block_tables=torch.zeros(1, 1, dtype=torch.int32), ++ seq_lens=torch.tensor([64], dtype=torch.int32), ++ max_seq_len=64, ++ softmax_scale=1.0, ++ **kwargs, ++ ) ++ ++ ++def test_cute_dsl_mla_decode_padded_q_tile_via_public_api(): ++ """Exercise the reported H96/Sq8 shape through the public dispatcher.""" ++ _run_padded_q_tile_case( ++ batch_size=1, ++ seq_len_k=128, ++ num_heads=96, ++ q_len=8, ++ dtype=torch.bfloat16, ++ enable_pdl=False, ++ via_public_api=True, ++ ) ++ ++ ++def test_compute_q_tile_layout(): ++ """Unit test the shared host/kernel query-tile geometry.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ from flashinfer.cute_dsl.attention.monolithic.mla_helpers import ( ++ compute_q_tile_layout, + ) + + cases = [ +- # (num_heads, seq_len_q, m_tile, expected) +- (128, 1, 128, 1), # H == M_tile → no fold +- (128, 4, 128, 1), # H == M_tile → no fold +- (64, 1, 128, 1), # seq_len_q=1 → F=1 +- (64, 2, 128, 2), # exact divisor, H*F=128 ≤ M_tile +- (64, 4, 128, 2), # H*F ≤ 128 caps F at 2; 4 % 2 == 0 +- (64, 3, 128, 1), # 3's only divisors are 1 and 3; H*3=192 > M_tile → F=1 +- (32, 4, 128, 4), # tighter pack: F=4, H*F=128 +- (32, 8, 128, 4), # capped by M_tile/H = 4 +- (32, 3, 128, 3), # max_fold=min(3, 4)=3; 3 % 3 == 0 → F=3 +- (16, 8, 128, 8), # max_fold=min(8, 8)=8; 8 % 8 == 0 → F=8 +- (16, 6, 128, 6), # max_fold=min(6, 8)=6; 6 % 6 == 0 → F=6 ++ # (H, Sq, M, (total_rows, num_tiles, tail_rows)) ++ (128, 3, 128, (384, 3, 128)), ++ (96, 1, 128, (96, 1, 96)), ++ (96, 3, 128, (288, 3, 32)), ++ (48, 8, 128, (384, 3, 128)), ++ (12, 11, 128, (132, 2, 4)), + ] + for H, S_q, m_tile, expected in cases: +- assert FP16.compute_fold_sq_ratio(H, S_q, m_tile) == expected, ( +- f"FP16.compute_fold_sq_ratio({H}, {S_q}, {m_tile}) " +- f"= {FP16.compute_fold_sq_ratio(H, S_q, m_tile)}, expected {expected}" ++ assert compute_q_tile_layout(H, S_q, m_tile) == expected, ( ++ f"compute_q_tile_layout({H}, {S_q}, {m_tile}) " ++ f"= {compute_q_tile_layout(H, S_q, m_tile)}, expected {expected}" + ) +- assert FP8.compute_fold_sq_ratio(H, S_q, m_tile) == expected, ( +- f"FP8.compute_fold_sq_ratio({H}, {S_q}, {m_tile}) " +- f"= {FP8.compute_fold_sq_ratio(H, S_q, m_tile)}, expected {expected}" ++ ++ for invalid in [(0, 1, 128), (129, 1, 128), (64, 0, 128), (64, 1, 0)]: ++ with pytest.raises(ValueError): ++ compute_q_tile_layout(*invalid) ++ ++ ++def test_nonpersistent_grid_y_limit(): ++ """Reject only nonpersistent grids beyond CUDA's Y-dimension limit.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _validate_nonpersistent_grid_y, ++ ) ++ ++ _validate_nonpersistent_grid_y(8_191, 8, is_persistent=False) ++ _validate_nonpersistent_grid_y(65_535, 1, is_persistent=False) ++ _validate_nonpersistent_grid_y(8_192, 8, is_persistent=True) ++ ++ with pytest.raises(ValueError, match=r"grid\.y would be 65536"): ++ _validate_nonpersistent_grid_y(8_192, 8, is_persistent=False) ++ ++ ++def test_mla_reducer_d_tile_selection(): ++ """Use output bands only when they shorten an underfilled reducer wave.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _get_reducer_d_tiles, ++ ) ++ ++ # B1/H32 and B1/H96 use D4; B1/H64 uses D2. H128 already fills a wave. ++ assert _get_reducer_d_tiles(1, 1, 32, 148, 32) == 4 ++ assert _get_reducer_d_tiles(1, 1, 64, 148, 32) == 2 ++ assert _get_reducer_d_tiles(1, 1, 96, 148, 32) == 4 ++ # Prefer the smaller tied topology and avoid duplication once rows fill a wave. ++ assert _get_reducer_d_tiles(1, 1, 48, 148, 32) == 2 ++ assert _get_reducer_d_tiles(1, 1, 128, 148, 32) == 1 ++ assert _get_reducer_d_tiles(4, 1, 96, 148, 32) == 1 ++ assert _get_reducer_d_tiles(1, 1, 96, 0, 32) == 1 ++ # Do not duplicate LSE work when a short sequence exposes too few splits. ++ assert _get_reducer_d_tiles(1, 1, 96, 148, 1) == 1 ++ assert _get_reducer_d_tiles(1, 1, 24, 148, 2) == 2 ++ ++ ++def test_mla_reducer_direct_class_capacity_defaults(): ++ """Direct class users keep the generic capacity unless opting into a cap.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ ++ import cutlass ++ ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp16 import ( ++ BlackwellMultiHeadLatentAttentionForwardFP16, ++ ) ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode_fp8 import ( ++ BlackwellMultiHeadLatentAttentionForwardFP8, ++ ) ++ from flashinfer.cute_dsl.attention.monolithic.mla_helpers import MAX_SPLITS ++ ++ kwargs = dict( ++ acc_dtype=cutlass.Float32, ++ lse_dtype=cutlass.Float32, ++ mma_qk_tiler_mn=(128, 128), ++ mma_pv_tiler_mn=(128, 256), ++ max_active_clusters=1, ++ page_size=64, ++ skip_correction_threshold=0.0, ++ is_persistent=True, ++ is_var_seq=False, ++ is_var_split_kv=False, ++ enable_pdl=False, ++ ) ++ for kernel_cls in ( ++ BlackwellMultiHeadLatentAttentionForwardFP16, ++ BlackwellMultiHeadLatentAttentionForwardFP8, ++ ): ++ assert kernel_cls(**kwargs).reducer_max_splits == MAX_SPLITS ++ assert kernel_cls(**kwargs, reducer_max_splits=32).reducer_max_splits == 32 ++ with pytest.raises(ValueError, match="variable split-KV"): ++ kernel_cls(**{**kwargs, "is_var_split_kv": True}, reducer_max_splits=32) ++ ++ ++def test_flat_q_tile_split_workspace_geometry(): ++ """Workspace sizing uses the flattened M128 query-tile count.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _get_split_kv_and_workspace_size, ++ ) ++ ++ batch_size, num_heads, q_len, num_q_tiles = 1, 48, 8, 3 ++ latent_dim = 512 ++ max_active_blocks = 148 ++ split_kv, workspace_size = _get_split_kv_and_workspace_size( ++ batch_size, ++ q_len, ++ num_heads, ++ latent_dim, ++ max_active_blocks, ++ ) ++ expected_split = min(max_active_blocks // (batch_size * num_q_tiles * 2), 32) ++ expected_workspace = ( ++ batch_size * 128 * num_q_tiles * expected_split * (latent_dim + 1) * 4 ++ ) ++ assert split_kv == expected_split ++ assert workspace_size == expected_workspace ++ ++ # Once one cluster per (batch, query tile) fills the machine, split-KV is ++ # disabled and no workspace is required. ++ direct_batch = max_active_blocks // (num_q_tiles * 2) ++ split_kv, workspace_size = _get_split_kv_and_workspace_size( ++ direct_batch, ++ q_len, ++ num_heads, ++ latent_dim, ++ max_active_blocks, ++ ) ++ assert split_kv == 1 ++ assert workspace_size == 0 ++ ++ ++def test_h96_sq8_split_workspace_drops_empty_partition(): ++ """Size only the eight nonempty K partitions at B1/H96/Sq8.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _get_split_kv_and_workspace_size, ++ ) ++ ++ split_kv, workspace_size = _get_split_kv_and_workspace_size( ++ 1, 8, 96, 512, 148, 1024 ++ ) ++ assert split_kv == 8 ++ assert workspace_size == 1 * 128 * 6 * 8 * (512 + 1) * 4 ++ ++ ++def test_cute_dsl_workspace_sizer_follows_selected_impl(): ++ """Autotuning must size workspace with the implementation it will launch.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ ++ from flashinfer.cute_dsl.attention.monolithic.mla_decode import ( ++ _get_split_kv_and_workspace_size as monolithic_sizer, ++ ) ++ from flashinfer.cute_dsl.attention.wrappers.batch_mla import ( ++ _get_split_kv_and_workspace_size as modular_sizer, ++ ) ++ from flashinfer.mla._core import _get_cute_dsl_workspace_sizer ++ ++ assert _get_cute_dsl_workspace_sizer("monolithic", None) is monolithic_sizer ++ assert _get_cute_dsl_workspace_sizer("modular", None) is modular_sizer ++ assert _get_cute_dsl_workspace_sizer("auto", None) is monolithic_sizer ++ assert _get_cute_dsl_workspace_sizer("auto", torch.empty(0)) is modular_sizer ++ ++ ++def test_monolithic_workspace_cap_drops_empty_partitions(): ++ """Autotuning must use the launcher's max-sequence split normalization.""" ++ if not is_cute_dsl_available(): ++ pytest.skip("CuTe DSL not available") ++ ++ from flashinfer.mla._core import _cute_dsl_max_supported_batch ++ ++ # One K tile needs no split workspace for any candidate batch. Without ++ # max_seq_len propagation this is conservatively sized as 32 splits and ++ # the zero-byte workspace incorrectly caps the autotune sweep at B=1. ++ assert ( ++ _cute_dsl_max_supported_batch( ++ workspace_bytes=0, ++ q_len=1, ++ num_heads=128, ++ kv_lora_rank=512, ++ max_active_blocks=148, ++ max_seq_len=128, ++ candidate_max=8, ++ cute_dsl_impl="monolithic", ++ sinks=None, + ) ++ == 8 ++ ) + + + @pytest.mark.parametrize("batch_size", [1, 4, 16]) +diff --git a/tests/autotuner/test_autotuner_core.py b/tests/autotuner/test_autotuner_core.py +index a36c1e88..a5cf0100 100644 +--- a/tests/autotuner/test_autotuner_core.py ++++ b/tests/autotuner/test_autotuner_core.py +@@ -14,6 +14,7 @@ from flashinfer.fused_moe.utils import ( + make_hybrid_bucket_mapper, + ) + from flashinfer.mla._core import ( ++ CuteDslMlaDecodeRunner, + _build_mla_decode_tuning_config, + _mla_decode_tuning_config, + ) +@@ -1446,3 +1447,34 @@ def test_find_nearest_profile_cache_dedups_mla_decode_config(): + finally: + AutoTuner._find_nearest_profile.cache_clear() + _mla_decode_tuning_config.cache_clear() ++ ++ ++def _cute_dsl_runner_cache_extras(max_seq_len: int, workspace_bytes: int): ++ runner = object.__new__(CuteDslMlaDecodeRunner) ++ runner.kv_cache = torch.empty((1, 32, 576), dtype=torch.bfloat16) ++ runner.workspace_buffer = torch.empty(workspace_bytes, dtype=torch.uint8) ++ runner.qk_nope_head_dim = 512 ++ runner.kv_lora_rank = 512 ++ runner.qk_rope_head_dim = 64 ++ runner.page_size = 32 ++ runner.max_seq_len = max_seq_len ++ runner.is_var_seq = True ++ runner.uses_shared_paged_kv_idx = True ++ runner.enable_pdl = False ++ runner.sinks = None ++ runner.cute_dsl_impl = "auto" ++ runner._resolved_cute_dsl_impl = "monolithic" ++ ++ query = torch.empty((1, 1, 128, 576), dtype=torch.bfloat16) ++ out = torch.empty((1, 1, 128, 512), dtype=torch.bfloat16) ++ return runner.get_cache_key_extras([query, None, None, out]) ++ ++ ++def test_cute_dsl_runner_cache_tracks_split_workspace_geometry(): ++ """Cache hits must not bypass sequence- or capacity-dependent validity.""" ++ key_257 = _cute_dsl_runner_cache_extras(257, 1_000_000) ++ key_385 = _cute_dsl_runner_cache_extras(385, 1_000_000) ++ assert key_257 != key_385 ++ ++ key_small_workspace = _cute_dsl_runner_cache_extras(257, 800_000) ++ assert key_257 != key_small_workspace diff --git a/docker/kimi_k3/kimi_k3_cu12.Dockerfile b/docker/kimi_k3/kimi_k3_cu12.Dockerfile new file mode 100644 index 000000000..c028a0bb1 --- /dev/null +++ b/docker/kimi_k3/kimi_k3_cu12.Dockerfile @@ -0,0 +1,136 @@ +# Kimi-K3 serving image (x86_64 / CUDA 12.9 / sm_90 + sm_100a). +# +# Base ships stock SGLang (editable at /sgl-workspace/sglang), DeepEP source +# (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package, +# and the CUDA 12.9 toolchain. +# +# This image adds the four Kimi-K3-specific pieces that stock lacks: +# 1. the Kimi-K3 SGLang code (this repo), editable-installed +# 2. DeepEP patch + rebuild: +# topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment, +# and cross-node timeout headroom; rebuilt for sm_90 and sm_100a only +# 3. DeepGEMM mega-MoE SiTU patch: +# JIT-header sentinel (activation_clamp==0.03125 -> K3 SiTU); no rebuild +# 4. FlashInfer CuTeDSL MLA DCP patch: +# apply the seven runtime-file diffs; exclude tests absent from the wheel +# +# Build (on/for x86_64; nvcc cross-compiles the DeepEP cubin, no GPU needed): +# docker build -f docker/kimi_k3/kimi_k3_cu12.Dockerfile \ +# --build-arg 'TORCH_CUDA_ARCH_LIST=9.0;10.0a' -t kimi-k3-cu129 . +# +# The FlashInfer MXFP4 MoE runner cubins are installed in the image below. +# The runner is auto-selected on SM100; the remaining kernel sources +# JIT-compile from the installed FlashInfer wheel on first launch and are +# cached. + +FROM lmsysorg/sglang:v0.5.16-cu129 AS base + +# Current Kimi-K3 source auto-discovers and builds its PyO3 extensions. +ARG RUST_VERSION="1.90.0" +ENV RUSTUP_HOME="/usr/local/rustup" \ + CARGO_HOME="/usr/local/cargo" \ + PATH="/usr/local/cargo/bin:${PATH}" +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + unzip \ + wget && \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ + sh -s -- -y --no-modify-path --profile minimal \ + --default-toolchain "${RUST_VERSION}" && \ + cargo --version && \ + rustc --version && \ + rm -rf /var/lib/apt/lists/* + +# Build one DeepEP wheel with native cubins for Hopper and B200. This CUDA 12.9 +# recipe intentionally excludes GB300 (sm_103/sm_103a). +ARG TORCH_CUDA_ARCH_LIST="9.0;10.0a" +RUN set -eu; \ + for arch in $(printf '%s' "${TORCH_CUDA_ARCH_LIST}" | tr ';' ' '); do \ + case "${arch}" in \ + 10.3|10.3a) \ + echo "ERROR: CUDA 12.9 image does not support SM103 (${arch})" >&2; \ + exit 1 \ + ;; \ + esac; \ + done + +# --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) --- +# Keep the installed extension modules, but discard Rust and pip build +# artifacts that are not used at runtime. +RUN rm -rf /sgl-workspace/sglang && \ + git clone --branch kimi-k3 \ + https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \ + cd /sgl-workspace/sglang && \ + rm -rf .git && \ + test ! -e .git && \ + pip install -e python --no-deps && \ + rm -rf \ + rust/target \ + rust/sglang-grpc/target \ + rust/sglang-mm/target \ + rust/sglang-server/target \ + /usr/local/cargo/registry \ + /root/.cache/pip + +# --- 2. DeepEP: patch (topk16 / hidden3584 / SourceMeta) + multi-arch rebuild --- +RUN TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + bash /sgl-workspace/sglang/docker/kimi_k3/apply_deepep_k3_patch.sh && \ + rm -rf /sgl-workspace/DeepEP/build /sgl-workspace/DeepEP/dist + +# --- 3. DeepGEMM mega-MoE: SiTU JIT-header patch (runtime-JIT, no rebuild) --- +RUN python3 /sgl-workspace/sglang/docker/kimi_k3/apply_deepgemm_situ_patch.py + +# Install the pinned FlashInfer MXFP4 MoE runner cubin pool. +ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip" +ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192" +ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1" +ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool" + +RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \ + cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \ + wget --no-verbose --output-document="${cubin_archive}" \ + "${TRTLLM_GEN_MOE_CUBIN_URL}" && \ + echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \ + sha256sum --check --strict - && \ + mkdir -p "${cubin_extract_dir}" && \ + unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \ + test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ + mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \ + "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ + test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \ + -type f -name '*.cubin' | wc -l)" -eq 1696 && \ + rm -f "${cubin_archive}" && \ + rm -rf "${cubin_extract_dir}" + +# Reinstall the matching FlashInfer package trio before patching its Python +# sources. A mixed Python/cubin/JIT-cache installation fails at import time. +# flashinfer-python and flashinfer-cubin are CUDA-independent packages; the +# JIT-cache wheel is selected from the official CUDA 12.9 index. +RUN python3 -m pip uninstall -y \ + flashinfer-python flashinfer-cubin flashinfer-jit-cache && \ + rm -rf /root/.cache/flashinfer /root/.cache/pip && \ + python3 -m pip install --no-deps \ + "flashinfer-python==0.6.15.post1" && \ + python3 -m pip install --no-deps \ + "flashinfer-cubin==0.6.15.post1" \ + --index-url https://flashinfer.ai/whl && \ + python3 -m pip install --no-deps \ + "flashinfer-jit-cache==0.6.15.post1" \ + --index-url https://flashinfer.ai/whl/cu129 && \ + python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; assert version("flashinfer-python").split("+", 1)[0] == expected; assert version("flashinfer-cubin").split("+", 1)[0] == expected; assert version("flashinfer-jit-cache").startswith(expected + "+cu129"), version("flashinfer-jit-cache")' && \ + rm -rf /root/.cache/pip + +ENV FLASHINFER_VERSION="0.6.15.post1" + +# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch --- +RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt && \ + FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \ + sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ + patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ + sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ + patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ + rm -rf /root/.cache/flashinfer /root/.cache/pip + +WORKDIR /sgl-workspace/sglang diff --git a/docker/kimi_k3/kimi_k3_cu13.Dockerfile b/docker/kimi_k3/kimi_k3_cu13.Dockerfile new file mode 100644 index 000000000..5b884cffa --- /dev/null +++ b/docker/kimi_k3/kimi_k3_cu13.Dockerfile @@ -0,0 +1,125 @@ +# Kimi-K3 serving image (aarch64 / sm_90 + sm_100a + sm_103a). +# +# Base ships stock SGLang (editable at /sgl-workspace/sglang), DeepEP source +# (deepseek-ai@d28bd67 at /sgl-workspace/DeepEP), the deep_gemm pip package, +# and the CUDA 13 toolchain (nvcc + /usr/local/cuda/include/cccl). +# +# This image adds the three Kimi-K3-specific pieces that stock lacks: +# 1. the Kimi-K3 SGLang code (this repo), editable-installed +# 2. DeepEP patch + rebuild: +# topk 11->16, SWITCH_HIDDEN += 3584, EP>8 SourceMeta alignment, +# cross-node timeout headroom, CUDA-13 cccl include; rebuilt for +# sm_90, sm_100a, and sm_103a +# 3. DeepGEMM mega-MoE SiTU patch: +# JIT-header sentinel (activation_clamp==0.03125 -> K3 SiTU); no rebuild +# 4. FlashInfer CuTeDSL MLA DCP patch: +# apply the seven runtime-file diffs; exclude tests absent from the wheel +# +# Build (on/for aarch64; nvcc cross-compiles the DeepEP cubin, no GPU needed): +# docker build -f docker/kimi_k3/kimi_k3_cu13.Dockerfile \ +# --build-arg 'TORCH_CUDA_ARCH_LIST=9.0;10.0a;10.3a' -t kimi-k3 . +# +# The FlashInfer MXFP4 MoE runner cubins are installed in the image below. +# The runner is auto-selected on SM100/103; the remaining kernel sources +# JIT-compile from the installed FlashInfer wheel on first launch and are +# cached. + +FROM lmsysorg/sglang:v0.5.16 AS base + +# Current Kimi-K3 source auto-discovers and builds its PyO3 extensions. +ARG RUST_VERSION="1.90.0" +ENV RUSTUP_HOME="/usr/local/rustup" \ + CARGO_HOME="/usr/local/cargo" \ + PATH="/usr/local/cargo/bin:${PATH}" +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + unzip \ + wget && \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | \ + sh -s -- -y --no-modify-path --profile minimal \ + --default-toolchain "${RUST_VERSION}" && \ + cargo --version && \ + rustc --version && \ + rm -rf /var/lib/apt/lists/* + +# Build one DeepEP wheel with native cubins for Hopper, B200, and GB300. +ARG TORCH_CUDA_ARCH_LIST="9.0;10.0a;10.3a" + +# --- 1. Kimi-K3 SGLang code (replaces the base's stock sglang, editable) --- +# Keep the installed extension modules, but discard Rust and pip build +# artifacts that are not used at runtime. +RUN rm -rf /sgl-workspace/sglang && \ + git clone --branch kimi-k3 \ + https://github.com/sgl-project/sglang.git /sgl-workspace/sglang && \ + cd /sgl-workspace/sglang && \ + rm -rf .git && \ + test ! -e .git && \ + pip install -e python --no-deps && \ + rm -rf \ + rust/target \ + rust/sglang-grpc/target \ + rust/sglang-mm/target \ + rust/sglang-server/target \ + /usr/local/cargo/registry \ + /root/.cache/pip + +# --- 2. DeepEP: patch (topk16 / hidden3584 / SourceMeta / cccl) + multi-arch rebuild --- +RUN TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" \ + bash /sgl-workspace/sglang/docker/kimi_k3/apply_deepep_k3_patch.sh && \ + rm -rf /sgl-workspace/DeepEP/build /sgl-workspace/DeepEP/dist + +# --- 3. DeepGEMM mega-MoE: SiTU JIT-header patch (runtime-JIT, no rebuild) --- +RUN python3 /sgl-workspace/sglang/docker/kimi_k3/apply_deepgemm_situ_patch.py + +# Install the pinned FlashInfer MXFP4 MoE runner cubin pool. +ARG TRTLLM_GEN_MOE_CUBIN_URL="https://github.com/sgl-project/whl/releases/download/trtllm_gen_moe_cubin_20260617/trtllm_gen_moe_cubin_pool_20260617_v0613rc1.zip" +ARG TRTLLM_GEN_MOE_CUBIN_SHA256="4900501cbe782a76b08a5858f9f07152287b97cb68114466dac286366b66c192" +ARG TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT="trtllm_gen_moe_cubin_pool_20260617_v0613rc1" +ENV SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL="/opt/trtllm_gen_moe_cubin_pool" + +RUN cubin_archive="/tmp/trtllm_gen_moe_cubin_pool.zip" && \ + cubin_extract_dir="/tmp/trtllm_gen_moe_cubin_extract" && \ + wget --no-verbose --output-document="${cubin_archive}" \ + "${TRTLLM_GEN_MOE_CUBIN_URL}" && \ + echo "${TRTLLM_GEN_MOE_CUBIN_SHA256} ${cubin_archive}" | \ + sha256sum --check --strict - && \ + mkdir -p "${cubin_extract_dir}" && \ + unzip -q "${cubin_archive}" -d "${cubin_extract_dir}" && \ + test ! -e "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ + mv "${cubin_extract_dir}/${TRTLLM_GEN_MOE_CUBIN_ARCHIVE_ROOT}" \ + "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" && \ + test "$(find "${SGLANG_TRTLLM_GEN_MOE_CUBIN_POOL}" \ + -type f -name '*.cubin' | wc -l)" -eq 1696 && \ + rm -f "${cubin_archive}" && \ + rm -rf "${cubin_extract_dir}" + +# Reinstall the matching FlashInfer package trio before patching its Python +# sources. A mixed Python/cubin/JIT-cache installation fails at import time. +RUN python3 -m pip uninstall -y \ + flashinfer-python flashinfer-cubin flashinfer-jit-cache && \ + rm -rf /root/.cache/flashinfer /root/.cache/pip && \ + python3 -m pip install --no-deps \ + "flashinfer-python==0.6.15.post1" && \ + python3 -m pip install --no-deps \ + "flashinfer-cubin==0.6.15.post1" \ + --index-url https://flashinfer.ai/whl && \ + python3 -m pip install --no-deps \ + "flashinfer-jit-cache==0.6.15.post1" \ + --index-url https://flashinfer.ai/whl/cu130 && \ + python3 -c 'from importlib.metadata import version; expected = "0.6.15.post1"; packages = ("flashinfer-python", "flashinfer-cubin", "flashinfer-jit-cache"); actual = {package: version(package).split("+", 1)[0] for package in packages}; assert all(value == expected for value in actual.values()), actual' && \ + rm -rf /root/.cache/pip + +ENV FLASHINFER_VERSION="0.6.15.post1" + +# --- 4. FlashInfer: CuTeDSL MLA decode-context-parallel runtime patch --- +RUN FLASHINFER_DCP_PATCH=/sgl-workspace/sglang/docker/kimi_k3/flashinfer-perkz-dcp-0.6.15.txt && \ + FLASHINFER_SITE_PACKAGES="$(python3 -c 'from pathlib import Path; import flashinfer; print(Path(flashinfer.__file__).resolve().parent.parent)')" && \ + sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ + patch --dry-run --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ + sed '/^diff --git a\/tests\//,$d' "${FLASHINFER_DCP_PATCH}" | \ + patch --batch --forward --strip=1 --directory="${FLASHINFER_SITE_PACKAGES}" && \ + rm -rf /root/.cache/flashinfer /root/.cache/pip + +WORKDIR /sgl-workspace/sglang diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 76348ec6f..ef29fa124 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -37,7 +37,7 @@ ENV BUILD_TRITON="0" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" -ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02" +ENV AITER_COMMIT_DEFAULT="v0.1.19" # =============================== # Base image 942 with rocm720 and args @@ -47,7 +47,7 @@ ENV BUILD_TRITON="1" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" -ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02" +ENV AITER_COMMIT_DEFAULT="v0.1.19" # =============================== # Base image 950 and args @@ -57,7 +57,7 @@ ENV BUILD_TRITON="0" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" -ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02" +ENV AITER_COMMIT_DEFAULT="v0.1.19" # =============================== # Base image 950 with rocm720 and args @@ -67,7 +67,7 @@ ENV BUILD_TRITON="1" ENV BUILD_LLVM="0" ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" -ENV AITER_COMMIT_DEFAULT="9127c94a18e4398e1eba91f6639e910f0994ad02" +ENV AITER_COMMIT_DEFAULT="v0.1.19" # Local source stage: with BRANCH_TYPE=local the build context is copied here and # used instead of git clone (mirrors docker/Dockerfile's local_src stage).