[Fix] Skip grid-based stage fallback for OTHER/MEMORY kernels to prevent small scheduling ops from misclassifying decode layers as prefill

[yichiche]

Summary: Grid-Based Stage Fallback Misclassification Fix

Problem

The 0303 profiling trace had severely degraded evaluation scores compared to the 0302 baseline:

┌──────────────────────┬─────────┬───────────────────┬──────────────────┐
│ Metric │ 0302 │ 0303 (before fix) │ 0303 (after fix) │
├──────────────────────┼─────────┼───────────────────┼──────────────────┤
│ S3 Round Consistency │ 100 (A) │ 44.4 (F) │ 100 (A) │
├──────────────────────┼─────────┼───────────────────┼──────────────────┤
│ S4 Type Sequence │ 100 (A) │ 57.3 (F) │ 100 (A) │
├──────────────────────┼─────────┼───────────────────┼──────────────────┤
│ Overall │ 100 (A) │ 75.4 (B) │ 100 (A) │
├──────────────────────┼─────────┼───────────────────┼──────────────────┤
│ Prefill layers │ 525 │ 1,157 │ 521 │
├──────────────────────┼─────────┼───────────────────┼──────────────────┤
│ Decode rounds │ 400 │ 798 │ 400 │
└──────────────────────┴─────────┴───────────────────┴──────────────────┘

Root Cause

Small scheduling/utility kernels (index_elementwise_kernel, CUDAFunctorOnSelf, arange) were being misclassified as PREFILL by the grid-based fallback rule (grid[0] < 200 → PREFILL). These kernels always have tiny grid dimensions regardless of whether
they run during prefill or decode.

In a typical decode layer, stage voting is razor-thin: 1 DECODE vote (from mla_a8w8_qh16_qseqlen1) vs 25 UNKNOWN. When 1-3 of these small kernels appeared as extras in a decode layer, their false PREFILL votes (via grid fallback) outnumbered the single
DECODE vote, flipping the layer to PREFILL. This broke decode rounds into fragments, cascading through stage propagation to corrupt adjacent tie-breaking layers.

Fix

profile/trace_analyzer.py — two changes:

1. Added kernel_type parameter to classify_stage() and skipped the grid-based fallback for OTHER and MEMORY kernel types, since these utility ops have small grids regardless of execution phase.
2. Reordered the call site to classify kernel type before stage, so the type can be passed in.

The fix is minimal (6 lines changed in the classifier) and preserves the grid fallback for compute-relevant kernels (attention, quantization, linear, MoE) where grid size is a meaningful prefill/decode signal.

Also Included

• /profile skill (.claude/commands/profile.md + profile/profile.md) — new slash command documenting the profiling toolchain, output file formats, analysis workflow, and tool invocation commands.