diff --git a/docs/api/comm.rst b/docs/api/comm.rst index 78f6bb4ce5b..801bc6ca22a 100644 --- a/docs/api/comm.rst +++ b/docs/api/comm.rst @@ -773,6 +773,40 @@ Prepared QKV workspaces and ordinary staging workspaces are not interchangeable. .. autoclass:: UlyssesQKVWorkspace :members: +Low-Precision Implementation Primitives +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The following internal helpers support numerical tests, trace definitions, +and kernel benchmarks. They are not exported from ``flashinfer.comm``; +applications should use the prepared ``UlyssesCommunicator`` QKV API above. + +.. currentmodule:: flashinfer.comm._ulysses_lowp + +.. autosummary:: + :toctree: ../generated + + capability + payload_spec + k_sum_v_amax + q_grouped_amax + k_grouped_amax + boundary_descriptors + merge_boundary_amax + k_boundary_minmax + derive_k_boundary_amax + zero_scale_and_padding + quant_q_into_payload + quant_kv_into_payload + quant_qkv_pack + unpack_for_sage + local_stats + finalize_stats + quant_and_pack + verify_duplicate_scale_slots + quant_v_fp8_with_scale + +.. currentmodule:: flashinfer.comm + Head-Chunk Layout and Transport Primitives ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/docs/api/gemm.rst b/docs/api/gemm.rst index a1715f7d392..2dad5c46c0b 100644 --- a/docs/api/gemm.rst +++ b/docs/api/gemm.rst @@ -41,7 +41,7 @@ Quantized-base GEMM with runtime BF16 LoRA (SM90, SM100, SM103) BF16 LoRA factors consumed as given (Cake generated programs). ``lora_scale=None`` disables the adapter even if ``d`` / ``lora_b`` are passed; pass a device scalar to apply it. The NVFP4 operator also takes a prepared adapter (``lora_adapter=prepare_nvfp4_lora_adapter(lora_b, -lora_scale, alpha)``: the adapter folded once at the switch, SVDQuant-equivalent numerics). +lora_scale, alpha)``: the adapter is folded once at the switch, with SVDQuant-equivalent numerics). .. autosummary:: :toctree: ../generated diff --git a/docs/api/kda.rst b/docs/api/kda.rst index f07774bf923..224df779fed 100644 --- a/docs/api/kda.rst +++ b/docs/api/kda.rst @@ -26,7 +26,7 @@ backend of its own. These prefill routes leave decode and speculative-decode dispatch unchanged. Automatic prefill policy on GB200 and GB300 -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ For one packed sequence, BF16 activations, head dimension 128 and an existing BF16 or FP32 state, ``auto`` uses these coarse performance regions: @@ -114,7 +114,7 @@ and decode retains its existing layout contract. recurrent_kda Static PTX prefill on B300 -~~~~~~~~~~~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~~~~~~~~~~~~ ``flashinfer.recurrent_kda(..., backend="ptx")`` explicitly selects four retained PTX programs from `NVlabs/kda (260927-kda-for-kda) @@ -715,7 +715,7 @@ should rotate fewer buffer sets, or call which releases all of it. Persistent CuTe DSL prefill with FP32 state -~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ``recurrent_kda(..., backend="cute-dsl-persistent")`` selects the persistent M128/BT32 kernel adapted from humanfia/kda-for-kda-release, now diff --git a/docs/api/mhc.rst b/docs/api/mhc.rst index a3e4f060a99..821b5d714cf 100644 --- a/docs/api/mhc.rst +++ b/docs/api/mhc.rst @@ -175,3 +175,19 @@ Names: ``flashinfer.cake_mhc_det_train.cake_mhc_det_train_`` -> non-deterministic package's nodes). The legacy module ``flashinfer.cake_mhc_train`` keeps the same ``cake_mhc_train_*`` names as aliases of the deterministic objects, without a ``deterministic`` keyword. + +.. currentmodule:: flashinfer.cake_mhc_train + +.. autosummary:: + :toctree: ../generated + + cake_mhc_train_forward_prepare + cake_mhc_train_forward_stats + cake_mhc_train_forward_finish + cake_mhc_train_backward_prepare + cake_mhc_train_backward_finish + cake_mhc_train_prepare + cake_mhc_train_finish + cake_mhc_train_head_forward + cake_mhc_train_head_backward + cake_mhc_train_head diff --git a/docs/api/norm.rst b/docs/api/norm.rst index 84977109e8f..dda406628c2 100644 --- a/docs/api/norm.rst +++ b/docs/api/norm.rst @@ -78,7 +78,7 @@ The names the training RMSNorm had before the deterministic package took its call computes ``dx`` row by row and adds the ``dw`` sums of its CTAs into FP64 sums in the workspace, in the order the CTAs finish: ``dw`` is not guaranteed bitwise reproducible; ``dx`` is, and it is within one BF16 ulp of - max(|dx|, rms(dx)) of the deterministic backward's ``dx`` (not always bitwise + ``max(|dx|, rms(dx))`` of the deterministic backward's ``dx`` (not always bitwise equal to it). In the windows the backward's notes name, a twosum single kernel keeps each thread's ``dw`` column sums exact (error-free FP32 products, TwoSum compensation) before the FP64 additions, so its ``dw`` error does not grow diff --git a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py index b18dc640626..27942584349 100644 --- a/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py +++ b/flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py @@ -15311,7 +15311,7 @@ def supported_hidden_sizes(arch: str) -> tuple[int, ...]: def weight_dtypes(arch: str, hidden: int, kind: str, residual: bool) -> tuple[str, ...]: - """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships for on ``arch``.""" + """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships on ``arch``.""" routes = ROUTES.get(arch, {}) return tuple( diff --git a/flashinfer/cake_rmsnorm_train.py b/flashinfer/cake_rmsnorm_train.py index adebb700b35..d845197b7ee 100644 --- a/flashinfer/cake_rmsnorm_train.py +++ b/flashinfer/cake_rmsnorm_train.py @@ -40,7 +40,7 @@ CTAs claim row chunks, write ``dx`` and add their per-column ``dw`` sums into FP64 sums in the workspace in the order they finish; the CTA that finishes last rounds the sums into ``dw``. ``dx`` is bitwise reproducible and within - one BF16 ulp of max(|dx|, rms(dx)) of the deterministic backward's ``dx`` + one BF16 ulp of ``max(|dx|, rms(dx))`` of the deterministic backward's ``dx`` (not always bitwise equal to it); ``dw`` is not guaranteed bitwise reproducible (the order of the FP64 additions depends on scheduling). In the windows the backward docstring's notes name, a twosum single kernel keeps diff --git a/flashinfer/cake_rmsnorm_train_backend/__init__.py b/flashinfer/cake_rmsnorm_train_backend/__init__.py index 86c7fd248ee..a38f153857d 100644 --- a/flashinfer/cake_rmsnorm_train_backend/__init__.py +++ b/flashinfer/cake_rmsnorm_train_backend/__init__.py @@ -1190,7 +1190,7 @@ def cake_rmsnorm_train_backward( for it: in the windows that ship it, one kernel that writes ``dx`` and adds the ``dw`` sums of its CTAs into FP64 sums in the workspace, which the CTA that finishes last rounds into ``dw`` (``dx`` is bitwise - reproducible and within one BF16 ulp of max(|dx|, rms(dx)) of the + reproducible and within one BF16 ulp of ``max(|dx|, rms(dx))`` of the deterministic kernel's ``dx``, not always bitwise equal to it; the order of the FP64 additions depends on scheduling, so ``dw`` is not guaranteed bitwise diff --git a/flashinfer/cake_rmsnorm_train_backend/cake_jit.py b/flashinfer/cake_rmsnorm_train_backend/cake_jit.py index 8dcd8518d33..52be7eee5bd 100644 --- a/flashinfer/cake_rmsnorm_train_backend/cake_jit.py +++ b/flashinfer/cake_rmsnorm_train_backend/cake_jit.py @@ -15585,7 +15585,7 @@ def supported_hidden_sizes(arch: str) -> tuple[int, ...]: def weight_dtypes(arch: str, hidden: int, kind: str, residual: bool) -> tuple[str, ...]: - """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships for on ``arch``.""" + """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships on ``arch``.""" routes = ROUTES.get(arch, {}) return tuple( diff --git a/flashinfer/cake_rmsnorm_train_backend/cake_windows.py b/flashinfer/cake_rmsnorm_train_backend/cake_windows.py index 5394abbca3e..5a5934ffd83 100644 --- a/flashinfer/cake_rmsnorm_train_backend/cake_windows.py +++ b/flashinfer/cake_rmsnorm_train_backend/cake_windows.py @@ -1435,7 +1435,7 @@ ``y`` by at most one BF16 ulp, on up to about 32 % of the rows (measured on B200, GB300 and R200). A backward window (``deterministic=False``) runs one of two forms. The single kernel computes ``dx`` row by row (bitwise reproducible, within one BF16 ulp of -max(|dx|, rms(dx)) of the deterministic backward's ``dx``, not always bitwise equal to +``max(|dx|, rms(dx))`` of the deterministic backward's ``dx``, not always bitwise equal to it) and adds the ``dw`` sums of its CTAs into FP64 column sums in a scheduling-dependent order: its ``dw`` is not bitwise reproducible, even for a fixed ``T``. The top window of ``bwd_h256`` on sm_100a and sm_107a, on sm_103a of @@ -1501,7 +1501,7 @@ ``y`` by at most one BF16 ulp, on up to about 32 % of the rows (measured on B200, GB300 and R200). A backward window (``deterministic=False``) runs one of two forms. The single kernel computes ``dx`` row by row (bitwise reproducible, within one BF16 ulp of -max(|dx|, rms(dx)) of the deterministic backward's ``dx``, not always bitwise equal to +``max(|dx|, rms(dx))`` of the deterministic backward's ``dx``, not always bitwise equal to it) and adds the ``dw`` sums of its CTAs into FP64 column sums in a scheduling-dependent order: its ``dw`` is not bitwise reproducible, even for a fixed ``T``. The top window of ``bwd_h256`` on sm_100a and sm_107a, on sm_103a of @@ -1653,7 +1653,7 @@ ``y`` by at most one BF16 ulp, on up to about 32 % of the rows (measured on B200, GB300 and R200). A backward window (``deterministic=False``) runs one of two forms. The single kernel computes ``dx`` row by row (bitwise reproducible, within one BF16 ulp of -max(|dx|, rms(dx)) of the deterministic backward's ``dx``, not always bitwise equal to +``max(|dx|, rms(dx))`` of the deterministic backward's ``dx``, not always bitwise equal to it) and adds the ``dw`` sums of its CTAs into FP64 column sums in a scheduling-dependent order: its ``dw`` is not bitwise reproducible, even for a fixed ``T``. The top window of ``bwd_h256`` on sm_100a and sm_107a, on sm_103a of diff --git a/flashinfer/gemm/cake_quant_lora_gemm.py b/flashinfer/gemm/cake_quant_lora_gemm.py index 9f3f545da8e..1f92b0402ba 100644 --- a/flashinfer/gemm/cake_quant_lora_gemm.py +++ b/flashinfer/gemm/cake_quant_lora_gemm.py @@ -743,7 +743,7 @@ def _lora_carrier(record: dict, device: torch.device) -> torch.Tensor: class StaleAdapterError(ValueError): - """A prepared adapter's source tensor (``lora_b`` / ``lora_scale`` / ``alpha``) moved since its prepare (snapshot + """A prepared adapter's source tensor (``lora_b`` / ``lora_scale`` / ``alpha``) moved since its preparation (snapshot semantics): re-prepare (``prepare_nvfp4_lora_adapter(..., out=adapter)``) or capture the refold inside the graph.""" diff --git a/flashinfer/quantization/kernels/nvfp4_kv_materialize.py b/flashinfer/quantization/kernels/nvfp4_kv_materialize.py index c59783e4775..a90ae045d9b 100644 --- a/flashinfer/quantization/kernels/nvfp4_kv_materialize.py +++ b/flashinfer/quantization/kernels/nvfp4_kv_materialize.py @@ -722,7 +722,23 @@ def __init__( device: Union[str, int, torch.device] = "cuda", uses_shared_paged_kv_idx: bool = True, ) -> None: - """Check the layout and device, and compile or load the kernel.""" + """Check the layout and device, and compile or load the kernel. + + Parameters + ---------- + num_kv_heads : int + Number of KV heads per page. + page_size : int + Tokens per page. Must be divisible by 4. + head_dim : int + Head dimension. Must be divisible by 64. + device : Union[str, int, torch.device] + CUDA device that holds the tensors passed to :meth:`run`. + uses_shared_paged_kv_idx : bool + Whether K and V share page indices. ``block_tables`` is + ``[batch_size, max_num_pages]`` if so, else + ``[batch_size, 2, max_num_pages]``. + """ device = torch.device(device) if device.type != "cuda": raise ValueError(f"device must be a CUDA device, got {device}")