[Docs] Split slash-combined docstring parameters (#57141)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
82731931ae
commit
5f203baedb
+8
-2
@@ -1482,8 +1482,14 @@ def cutlass_w4a8_moe_mm(
|
||||
Cumulative token offsets
|
||||
problem_sizes:
|
||||
Per-expert (M, N, K) GEMM sizes used by the grouped GEMM launcher.
|
||||
a/b/c/group_scale_strides:
|
||||
Strides describing the memory layout of the input tensors.
|
||||
a_strides:
|
||||
Strides describing the memory layout of a_tensors.
|
||||
b_strides:
|
||||
Strides describing the memory layout of b_tensors.
|
||||
c_strides:
|
||||
Strides describing the memory layout of out_tensors.
|
||||
group_scale_strides:
|
||||
Strides describing the memory layout of b_group_scales.
|
||||
maybe_schedule:
|
||||
Optional override to choose a specific kernel or epilogue schedule.
|
||||
|
||||
|
||||
@@ -190,9 +190,11 @@ def candidate_blocks_to_sparse_indices(
|
||||
candidate_blocks: [rows, K] int32 request-local candidate block ids
|
||||
(-1 padded), in units of ``candidate_block_size`` positions. K
|
||||
must be a power of two (the in-kernel sort); rows may be strided.
|
||||
row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same
|
||||
row_ks: [rows] int32 per-row K range start; bounds are in the same
|
||||
(packed-workspace) coordinates the sparse kernel iterates over.
|
||||
Pass zeros for the paged path, whose blocks are context-relative.
|
||||
row_ke: [rows] int32 per-row K range end, in the same coordinates as
|
||||
``row_ks``.
|
||||
candidate_block_size: Positions per candidate block.
|
||||
sparse_block_kv: Positions per sparse block (8 or 16).
|
||||
out: Optional ``(sparse_indices, end)`` buffers to write into,
|
||||
@@ -322,12 +324,16 @@ def sparse_mqa_logits_prefill_chunk(
|
||||
k_scale: [total_kv] int32 packed UE8M0 K scales.
|
||||
weights: [rows, H] bf16 per-head weights; the sparse kernels take
|
||||
bf16 and do not fold the Q scale in.
|
||||
cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the
|
||||
packed workspace.
|
||||
cu_seqlen_ks: [rows] int32 per-token K start bounds in the packed
|
||||
workspace.
|
||||
cu_seqlen_ke: [rows] int32 per-token K end bounds in the packed
|
||||
workspace.
|
||||
candidate_blocks: [rows, K] int32 request-local candidate block ids.
|
||||
topk_indices: [rows, topk_tokens] output buffer.
|
||||
sparse_indices/end/col_indices: Caller-owned scratch, see
|
||||
`candidate_blocks_to_sparse_indices` and `sparse_topk_remap`.
|
||||
sparse_indices: Caller-owned scratch, see
|
||||
`candidate_blocks_to_sparse_indices`.
|
||||
end: Caller-owned scratch, see `candidate_blocks_to_sparse_indices`.
|
||||
col_indices: Caller-owned scratch, see `sparse_topk_remap`.
|
||||
kernel_metadata: DeepGEMM schedule from a previous call with the same
|
||||
candidates and bounds (i.e. another indexer layer in the same
|
||||
step). When given, the candidate expansion is skipped and
|
||||
@@ -413,7 +419,9 @@ def sparse_mqa_logits_paged_decode(
|
||||
candidate_blocks: [rows, K] int32 candidate block ids.
|
||||
topk_indices: [rows, topk_tokens] output buffer.
|
||||
row_ks: [rows] int32 zeros (paged blocks are context-relative).
|
||||
sparse_indices/end/col_indices: Caller-owned scratch.
|
||||
sparse_indices: Caller-owned scratch.
|
||||
end: Caller-owned scratch.
|
||||
col_indices: Caller-owned scratch.
|
||||
kernel_metadata: See `sparse_mqa_logits_prefill_chunk`.
|
||||
|
||||
Returns:
|
||||
|
||||
@@ -726,8 +726,10 @@ def get_sparse_mqa_logits_metadata(
|
||||
need a host-side sync to pick one.
|
||||
|
||||
Args:
|
||||
cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV
|
||||
cu_seqlen_ks: Per-row K range start bounds in the packed KV
|
||||
workspace, shape [num_q_tokens], dtype int32.
|
||||
cu_seqlen_ke: Per-row K range end bounds in the packed KV workspace,
|
||||
shape [num_q_tokens], dtype int32.
|
||||
num_kv_tokens: Total KV tokens in the packed workspace.
|
||||
sparse_kv_block_indices: Per-row candidate block ids, shape
|
||||
[num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's
|
||||
|
||||
Reference in New Issue
Block a user