[AI] flag_gems topk/weight_norm/chebyshev_v ignore input strides #173

Issue Type

• Bug Report

AI Agent Information

• Agent: Claude Code CLI
• Model: Claude Sonnet 4.6
• Session Context: Follow-up to [AI] flag_gems softmax family ignores input strides (silent wrong gradients on transposed inputs) #171 (softmax strided bug). A batch scan
  of the survey's strided profile found 20 candidate ops; strict re-testing
  with identical CPU/GPU inputs isolated three real gems stride bugs.

Summary

Three flag_gems kernels give wrong results on transposed (strided) inputs,
same defect class as the softmax family (#171): topk (max diff 1.22 on a
32x32 transpose), _weight_norm_interface (1.40), and
special_chebyshev_polynomial_v (NaN). Contiguous inputs are correct. All
three are used in real workloads (topk: sampling/beam search; weight_norm:
weight normalization layers). Two other survey candidates were investigated
and excluded: index_copy (test artifact -- CPU/GPU reference used
independent random tensors) and unfold_backward (native CUDA itself
disagrees with the CPU reference, a PyTorch platform difference, not a
torch_fl/gems defect).

Environment (for bug reports)

FLAGOS_USE_FLAGGEMS=1 GEMS_VENDOR=nvidia
# flag_gems source: 7fb49bad47116434961bfb2b912811716d383eaf

Reproduction

import torch
import torch_fl

x = torch.randn(32, 32, device="flagos:0")
xt = x.t()
g = torch.randn(32, device="flagos:0")

torch.topk(xt, 5, dim=1).values  # vs CPU: max diff 1.22 (wrong)
torch.ops.aten._weight_norm_interface(xt, g)[0]  # vs CPU: max diff 1.40 (wrong)
torch.special.chebyshev_polynomial_v(xt, torch.ones(1, device="flagos:0"))  # NaN

Expected vs Actual Behavior

Expected:
Correct results for any input layout, as the boxed CUDA kernels produce.

Actual:
Transposed inputs give wrong values / NaN silently. Contiguous inputs are
correct.

Root Cause Analysis

Same defect class as #171: the gems kernels index memory assuming contiguous
strides. Verified per-op that the boxed CUDA kernels (PyTorch native) handle
transposed inputs correctly, so rerouting is a complete fix.

Proposed Solution

Add the three ops to flaggems_runtime_broken (project precedent for
strided-sensitive kernels, same as #171) and regenerate.

Verification Plan

• Manual: all three ops match CPU on transposed and contiguous inputs
  after the fix (identical CPU/GPU tensors); unfold_backward/
  index_copy excluded with evidence
• Conf/kernel consistency: test_flaggems_conf_consistency.py 6/6
• CUDA CI after merge

Context & Investigation

• Batch-scanned the survey's 2d-f32-strided profile: 20 candidates.
• Re-tested each with identical CPU/GPU tensors; excluded index_copy (test
  artifact) and unfold_backward (native-CUDA-vs-CPU platform difference,
  verified on stock CUDA: max diff 0.76 without torch_fl).

Related Code Locations

• scripts/codegen_ops.py:2530 -- flaggems_runtime_broken "(4)" block
• flag_gems/ops/topk.py, flag_gems/ops/weight_norm.py,
  flag_gems/ops/special_chebyshev_polynomial_v.py (upstream kernels)

Checklist - AI Agents MUST Complete All

• I have provided complete environment information
• I have included a minimal, self-contained reproducer
• I have included full error output with traceback
• I have analyzed the root cause (not just symptoms)
• I have proposed a specific solution with implementation approach
• I have identified affected code locations with line numbers
• I have described how to verify the fix
• I have checked for duplicate issues
• All text is in English (required per CLAUDE.md)
• Code follows project conventions (checked existing code style)