Krusty84/triton-ascend-agent-dev-kit

# Assign Triton-Ascend Token Pools

## Goal

Copy each request's contiguous cache segment into the correct row and token range of a request-to-token table.

## Workflow

1. Launch one program per request.
2. Load kv_start, kv_end, and the destination request-pool row.
3. Sum end-start for all earlier requests to find this request's source offset in the concatenated cache.
4. Use int32 metadata when all lengths, offsets, and pool indices fit.
5. Copy kv_end-kv_start elements in fixed BLOCK_SIZE chunks with derived per-iteration offsets.

# Batch Triton-Ascend Token Reorder

## Goal

Gather several indexed token rows per program, assemble them in an on-chip buffer, and write one contiguous output tile.

## Workflow

1. Treat input and output as S by D and indices as a length-S integer vector.
2. Assign BLOCK_SIZE consecutive output rows to each program.
3. Load their indices as one vector.
4. Loop over valid rows, load each indexed D-element source row, and insert it into a BLOCK_SIZE by D UB tensor.
5. Store the assembled tile with row and column masks.

# Batch Triton-Ascend Token Unreorder

## Goal

Read a contiguous block of routed token rows once, extract rows in UB, and scatter each row to its indexed destination.

## Workflow

1. Treat input as S by D contiguous routed rows and indices as destination rows.
2. Assign BLOCK_SIZE consecutive input rows to each program.
3. Load a BLOCK_SIZE by D input tile and BLOCK_SIZE indices with masks.
4. Extract each valid row with extension.extract_slice.
5. Store the row at indices[row] * D with a column mask.

# Build Triton-Ascend Binned MoE Routing

## Goal

Map sorted token assignments into a tensor shaped experts by expert_capacity by hidden_size, truncate overflow per expert, and implement the inverse scatter and router-weight gradient.

## Workflow

1. Require bin_ids to be nondecreasing expert IDs aligned with indices.
2. Interpret bins[e] as the inclusive cumulative assignment count through expert e.
3. For gather, compute each assignment's offset within its expert and keep only offsets below expert_capacity.
4. Buffer rows for the current expert in UB; flush them when the expert ID changes inside a sub-block.
5. For scatter, load valid expert-capacity rows, store to unique assignment slots, optionally scale, and reduce TOP_K slots per token.
6. For wgrad, dot each kept expert row with the gradient of indices[i] // TOP_K.

# Build Triton-Ascend Fused Attention

## Goal

Compute softmax(QKᵀ * scale)V without materializing the full attention matrix. Support BNSD tensors shaped Z by H by N_CTX by HEAD_DIM.

## Workflow

1. Validate matching Q, K, and V shapes, compatible strides, supported dtype, and HEAD_DIM in {16, 32, 64, 128, 256}.
2. Choose BLOCK_M for query rows and BLOCK_N for key/value rows. Require N_CTX to be divisible by both in the compact implementation.
3. Map each task to a batch, head, and query-block tuple and create block pointers for Q, K, V, and output.
4. Load one Q block and stream K/V blocks through an online-softmax inner loop.
5. For causal attention, process preceding blocks and the diagonal block as separate stages; apply the triangular mask only to the diagonal stage.
6. Normalize the accumulated output by the running denominator and store it in the output dtype.

# Build Triton-Ascend Fused Softmax

## Goal

Compute softmax across the last dimension while reading each input row once and writing each output row once.

## Workflow

1. Interpret the input as n_rows by n_cols and preserve its row stride.
2. Choose BLOCK_SIZE as the next power of two greater than or equal to n_cols.
3. Let each program process rows separated by tl.num_programs(0).
4. Load a padded row with -inf outside n_cols.
5. Subtract the row maximum, exponentiate, reduce the sum, normalize, and store with the original mask.

# Build Triton-Ascend LayerNorm

## Goal

Implement the forward equation y = ((x - mean) / sqrt(var + eps)) * weight + bias over the final tensor dimension.

## Workflow

1. Reshape the input logically to M by N, where N is the normalized dimension.
2. Allocate output with the original shape and allocate optional mean and reciprocal-standard-deviation buffers as float32 vectors of length M.
3. Launch one program per row.
4. Accumulate the mean in float32 across masked BLOCK_SIZE tiles.
5. Make a second pass for variance, compute rstd, then make a third pass to normalize and apply weight and bias.

# Build Triton-Ascend MoE Gather and Scatter

## Goal

Implement the unbinned routing triplet over assignment IDs: gather token rows into routing order, scatter weighted rows back to tokens, and compute one router-weight gradient per assignment.

## Workflow

1. Define L = tokens * TOP_K and require indices to be a permutation of assignment IDs in [0, L).
2. Gather routed row i from token row indices[i] // TOP_K.
3. Scatter routed row i into assignment slot indices[i], optionally scale by weights[indices[i]], then reduce the TOP_K slots per token.
4. Compute wgrad[indices[i]] as the float32 dot product of routed row i and grads[indices[i] // TOP_K].
5. Split L approximately evenly across the target vector cores.
6. Process SUB_BLOCK_SIZE assignments and BLOCK_X feature columns at a time.

# Build Triton-Ascend Padded MoE Routing

## Goal

Store every expert's real assignments in a separately padded segment, preserve zero-filled gaps, and implement inverse scatter and per-assignment weight gradients.

## Workflow

1. Sort assignments by expert to obtain aligned indices and nondecreasing bin_ids.
2. Build bins from real cumulative counts.
3. Round each expert count to the required padding multiple and build padded_bins from cumulative padded counts.
4. Map sorted position i to offset_in_expert = i - bins[e-1] and padded row = padded_bins[e-1] + offset_in_expert.
5. Gather token indices[i] // TOP_K into that padded row.
6. Scatter real rows back through unique assignment IDs and compute wgrad with the same mapping.

# Build Triton-Ascend Small Matmul

## Goal

Compute output[A, C] = x[A, B] @ y[B, C] + bias[A, C] in one Triton program for small compile-time shapes.

## Workflow

1. Validate x as A by B, y as B by C, and bias as A by C.
2. Create row, reduction, and column offsets with tl.arange.
3. Form the flattened pointer grids through broadcasting.
4. Load both matrices and bias, call tl.dot, and store the result.
5. Launch exactly one program only while all tiles fit the target's on-chip resources.

# Build Triton-Ascend Vector Add

## Goal

Implement a standalone NPU kernel that computes elementwise addition for arbitrary tensor lengths. Reuse the same pattern for other one-dimensional elementwise operations.

## Workflow

1. Require both inputs to have the same shape, dtype, device, and supported layout.
2. Allocate the output with torch.empty_like and use output.numel() as the logical length.
3. Assign one contiguous block to each Triton program.
4. Mask every load and store so non-divisible tail elements never access invalid memory.
5. Derive the launch grid from the selected BLOCK_SIZE.

# Fuse Triton-Ascend Cat and Causal Conv1d

## Goal

Join convolution state and new tokens in UB, update the cached tail, and compute grouped causal Conv1d without materializing a global concatenation.

## Workflow

1. Preserve the reference contract for x, conv_state, depthwise weights, bias, activation, and state-row indices.
2. Load state and input from nonnegative contiguous addresses.
3. Transpose or reshape them so concatenation occurs along a contiguous flattened UB axis.
4. Insert state followed by new tokens with extension.insert_slice.
5. Derive the new state tail and convolution windows from the assembled tensor.
6. Map work to approximately the available vector-core count and loop over remaining feature tiles inside each program.

# Max-Autotune Triton-Ascend

## Goal

Expand a small set of kernel meta-parameter configurations into a controlled search over Ascend compiler options.

## Workflow

1. Verify the kernel and PyTorch reference before tuning.
2. Import max_autotune from triton.backends.ascend.runtime.
3. Provide a small base configs list containing only meaningful kernel meta-parameters.
4. Set key to runtime values whose changes may require a new winner.
5. Classify the kernel as vector, cube, or mix.
6. Add short explicit option lists only when the built-in search space is unsuitable.

# Optimize Triton-Ascend Discrete Loads

## Goal

Replace repeated indirect global reads with one contiguous global-to-UB transfer followed by on-chip gather.

## Workflow

1. Confirm the source domain length M is bounded and fits the UB budget.
2. Load x[0:M] contiguously into an on-chip vector.
3. Load the N indices and validate both lower and upper bounds.
4. Replace invalid indices with a safe value before tl.gather.
5. Gather along the source axis and store only valid outputs.

# Optimize Triton-Ascend Grouped Attention Access

## Goal

Preserve vectorized contiguous loads while gathering KV-cache rows selected by an indirect token-location vector.

## Workflow

1. Classify each target tile by which axis is discrete and which axis is contiguous.
2. Load a discrete outer row and its contiguous head dimension as one vector.
3. Assemble BLOCK_N rows into an on-chip tensor with extension.insert_slice.
4. Transpose the assembled K tile when tl.dot requires HEAD_DIM by BLOCK_N.
5. Load V directly as BLOCK_N by VALUE_DIM when its inner value dimension is contiguous.
6. Feed the tiles into numerically stable online softmax across KV splits.

# Optimize Triton-Ascend Integer Vectors

## Goal

Keep integer elementwise work on efficient Ascend vector paths without changing numerical semantics.

## Workflow

1. Determine the minimum signed range required by inputs, offsets, intermediates, and outputs.
2. Use torch.int32 tensors and int32-compatible scalar arguments when every possible value fits.
3. Keep pointer arithmetic and tail masking separate from the arithmetic dtype decision.
4. Launch a masked one-dimensional kernel and compare it with an int32 PyTorch reference.
5. Benchmark int32 and int64 only after synchronizing the NPU around timed regions.

# Prevent Triton-Ascend UB Overflow

## Goal

Keep each live tile within Ascend UB while preserving the semantics of a larger logical operation.

## Workflow

1. List all simultaneously live tensors and estimate their bytes, including padding and multibuffer copies.
2. Select a BLOCK_SIZE that fits the UB budget with safety margin.
3. Replace one load over the maximum logical length with a compile-time loop over BLOCK_SIZE tiles.
4. Recompute global offsets from the loop index and apply a runtime validity mask.
5. Store each completed tile before advancing so earlier tiles do not remain live.

# Rank Triton-Ascend Configs with Costmodel

## Goal

Predict latency for several compile-time kernel configurations and retain the most promising candidates for real correctness checks and benchmarking.

## Workflow

1. Define the Triton kernel, pointer/runtime signature, and candidate constexpr values.
2. Build a separate ASTSource and TTIR string for every candidate.
3. Enable the costmodel backend in Ascend compiler options.
4. Create one item per candidate with a unique config name, TTIR, and optional arg_bindings.
5. Call costmodel_bench, sort the returned microsecond predictions, and shortlist candidates.
6. Verify and benchmark the shortlist on real NPU inputs.

# Schedule Triton-Ascend Loads

## Goal

Issue independent memory work before a dependency-blocked load so the Ascend pipeline can overlap operations across loop iterations.

## Workflow

1. Draw read/write dependencies for every load and store in one loop iteration and across adjacent iterations.
2. Identify a load blocked by a prior store to the same or possibly aliased address.
3. Find later loads that have no dependency on that store.
4. Move only those independent loads before the blocked load.
5. Preserve calculation and store semantics, then inspect the generated schedule or profiler trace.

# Skip Triton-Ascend Padding Fill

## Goal

Avoid filling invalid load lanes when those values have no influence on any observable result.

## Workflow

1. Trace every consumer of the loaded tensor.
2. Prove that lanes where mask is false cannot reach arithmetic, reductions, indexing, control flow, or stores.
3. Omit other and set care_padding=False on the masked load.
4. Preserve downstream masks until the loaded value is no longer live.
5. Benchmark against the default padded load on the same runtime-tail shapes.

# Tile Triton-Ascend Gather

## Goal

Compute out[b, k] = x[b, index[b, k]] without launching one logical program for every batch-row and K tile.

## Workflow

1. Query the target num_vectorcore value.
2. Choose BLOCK_B so ceil(B / BLOCK_B) is close to, and not far above, the core count.
3. Launch a one-dimensional grid over batch blocks.
4. Loop over K inside each program in BLOCK_K tiles.
5. Load index tiles, validate them against C, gather x values with explicit strides, and store with combined masks.

# Tune Triton-Ascend Kernels

## Goal

Select a correct kernel configuration for each relevant input shape without hardcoding one block size for all workloads.

## Workflow

1. Establish a correct untuned kernel and reference test.
2. Choose standard autotune when candidate configurations are known.
3. Choose advanced autotune when Triton-Ascend should infer vector split or tiling candidates.
4. Put only values that can change the winner in key.
5. Derive the launch grid from meta-parameters and omit tunable meta-parameters at launch.
6. Validate the tuned result, then benchmark representative shapes.

# Validate Triton-Ascend Accuracy

## Goal

Make correctness checks explicit and consistent across Triton-Ascend kernel tests.

## Workflow

1. Compute the PyTorch reference from the same inputs before interpreting performance.
2. Assert identical output shape and dtype.
3. Compare floating-point tensors with dtype-aware tolerances.
4. Compare integer and boolean tensors exactly.
5. Add edge shapes, tail masks, zeros, extreme values, and NaNs when the operator semantics allow them.

# Vectorize Triton-Ascend Comparisons

## Goal

Move eligible explicit comparisons onto Ascend vector cast and compare instructions without changing mask semantics.

## Workflow

1. Identify comparisons used in compute expressions such as tl.where, not only load/store masks.
2. Prove that all compared integer values are exactly representable in float32.
3. Cast the index vector to tl.float32 immediately before the comparison.
4. Keep original integer offsets for pointer arithmetic and memory masks.
5. Inspect performance and compare results around the tail boundary.